""" model_loader.py --------------- Loads all AI models: - LLM: epfl-llm/meditron-7b (4-bit NF4, ~5GB VRAM) - Vision: Salesforce/blip2-opt-2.7b (float16, ~3.5GB VRAM) VRAM budget on RTX 4060 8GB: Meditron-7B 4-bit : ~5.0 GB BLIP-2 : ~3.5 GB NOTE: Run them sequentially (caption first, then LLM) so both never sit in VRAM simultaneously. Or use 8GB by loading BLIP-2, captioning, unloading, then loading Meditron. This is handled automatically. """ import gc import logging import os from pathlib import Path import torch logger = logging.getLogger(__name__) SLM_MODEL_ID = os.getenv("SLM_MODEL_ID", "BioMistral/BioMistral-7B") VISION_MODEL_ID = os.getenv("VISION_MODEL_ID", "Salesforce/blip2-opt-2.7b") LORA_ADAPTER_PATH = os.getenv("LORA_ADAPTER_PATH", "models/slm_model/lora_adapter") USE_4BIT = os.getenv("USE_4BIT", "true").lower() == "true" DEVICE = "cuda" if torch.cuda.is_available() else "cpu" logger.info(f"Device: {DEVICE} | 4-bit: {USE_4BIT}") if DEVICE == "cuda": logger.info( f"GPU: {torch.cuda.get_device_name(0)} | " f"VRAM: {torch.cuda.get_device_properties(0).total_memory/1e9:.1f}GB" ) _vision_model = None _vision_processor = None _slm_model = None _slm_tokenizer = None # ═══════════════════════════════════════════════════════════════════════════ # VISION — BLIP-2 (generates real image captions) # ═══════════════════════════════════════════════════════════════════════════ def load_vision_encoder(): global _vision_model, _vision_processor if _vision_model is not None: return _vision_model, _vision_processor logger.info(f"Loading BLIP-2: {VISION_MODEL_ID}") from transformers import Blip2ForConditionalGeneration, Blip2Processor _vision_processor = Blip2Processor.from_pretrained( VISION_MODEL_ID, cache_dir="models/vision_encoder" ) if DEVICE == "cuda": _vision_model = Blip2ForConditionalGeneration.from_pretrained( VISION_MODEL_ID, torch_dtype=torch.float16, device_map="auto", cache_dir="models/vision_encoder", ) else: _vision_model = Blip2ForConditionalGeneration.from_pretrained( VISION_MODEL_ID, torch_dtype=torch.float32, cache_dir="models/vision_encoder", ) _vision_model.eval() logger.info("BLIP-2 loaded.") return _vision_model, _vision_processor def unload_vision_encoder(): """Free VRAM after captioning so Meditron-7B can load.""" global _vision_model, _vision_processor if _vision_model is not None: del _vision_model _vision_model = None _vision_processor = None if DEVICE == "cuda": torch.cuda.empty_cache() gc.collect() logger.info("BLIP-2 unloaded — VRAM freed for Meditron.") # ═══════════════════════════════════════════════════════════════════════════ # LLM — Meditron-7B (medical language model) # ═══════════════════════════════════════════════════════════════════════════ def load_slm(): global _slm_model, _slm_tokenizer if _slm_model is not None: return _slm_model, _slm_tokenizer logger.info(f"Loading LLM: {SLM_MODEL_ID} | 4-bit: {USE_4BIT}") from transformers import AutoModelForCausalLM, AutoTokenizer _slm_tokenizer = AutoTokenizer.from_pretrained( SLM_MODEL_ID, trust_remote_code=True, padding_side="left", cache_dir="models/slm_model", ) if _slm_tokenizer.pad_token is None: _slm_tokenizer.pad_token = _slm_tokenizer.eos_token if USE_4BIT and DEVICE == "cuda": from transformers import BitsAndBytesConfig bnb = BitsAndBytesConfig( load_in_4bit=True, bnb_4bit_use_double_quant=True, bnb_4bit_quant_type="nf4", bnb_4bit_compute_dtype=torch.bfloat16, ) _slm_model = AutoModelForCausalLM.from_pretrained( SLM_MODEL_ID, quantization_config=bnb, device_map="auto", trust_remote_code=True, cache_dir="models/slm_model", ) logger.info("Meditron-7B loaded 4-bit NF4 (~5GB VRAM)") elif DEVICE == "cuda": _slm_model = AutoModelForCausalLM.from_pretrained( SLM_MODEL_ID, torch_dtype=torch.float16, device_map="auto", trust_remote_code=True, cache_dir="models/slm_model", ) logger.info("Meditron-7B loaded float16 (~14GB VRAM)") else: logger.warning("No GPU — CPU inference will be very slow") _slm_model = AutoModelForCausalLM.from_pretrained( SLM_MODEL_ID, torch_dtype=torch.float32, trust_remote_code=True, low_cpu_mem_usage=True, cache_dir="models/slm_model", ) _slm_model = _apply_lora_adapter(_slm_model) _slm_model.eval() logger.info("Meditron-7B ready.") return _slm_model, _slm_tokenizer def _apply_lora_adapter(model): adapter_path = Path(LORA_ADAPTER_PATH) if not (adapter_path / "adapter_config.json").exists(): logger.info("No LoRA adapter — using base Meditron-7B.") return model try: from peft import PeftModel model = PeftModel.from_pretrained(model, str(adapter_path)) model = model.merge_and_unload() logger.info(f"LoRA adapter loaded from {adapter_path}") except Exception as e: logger.warning(f"LoRA load failed: {e} — using base model.") return model # ═══════════════════════════════════════════════════════════════════════════ # Image captioning — BLIP-2 # ═══════════════════════════════════════════════════════════════════════════ def generate_image_caption(pil_image, vision_model=None, vision_processor=None) -> str: if vision_model is None or vision_processor is None: vision_model, vision_processor = load_vision_encoder() prompt = ( "Question: Describe this medical radiology image in detail. " "Include all visible anatomy, any hardware or implants, " "any pathology or abnormalities, and the image quality. Answer:" ) try: inputs = vision_processor( images=pil_image, text=prompt, return_tensors="pt" ) if DEVICE == "cuda": inputs = { k: v.to(DEVICE).to(torch.float16) if v.dtype == torch.float32 else v.to(DEVICE) for k, v in inputs.items() } with torch.no_grad(): ids = vision_model.generate( **inputs, max_new_tokens=150, num_beams=4, repetition_penalty=1.3, ) caption = vision_processor.decode(ids[0], skip_special_tokens=True) if "Answer:" in caption: caption = caption.split("Answer:")[-1].strip() logger.info(f"BLIP-2 caption: {caption[:150]}") return caption except Exception as e: logger.warning(f"BLIP-2 failed: {e}") return "" def extract_visual_features(pil_image, vision_model=None, vision_processor=None): caption = generate_image_caption(pil_image, vision_model, vision_processor) return None, caption