Spaces:
No application file
No application file
| """ | |
| 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 | |