model_loader.py

"""
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