models/model_manager.py

"""
Model Manager.

DESIGN DECISION: this entire application runs on free CPU Basic hardware —
no ZeroGPU is requested anywhere. All three models were deliberately chosen
small enough (343MB classifier, ~100MB captioner, 0.5B-param LLM) that CPU
inference is fast enough for a live demo without fighting GPU quota/cold-start
behavior. This is the "smallest model that produces a convincing demo"
principle applied literally.

Models are loaded lazily (only when first needed) and cached as module-level
singletons so repeated calls during one demo run don't reload from disk.
Every load is wrapped so a failure degrades the relevant agent instead of
crashing the whole Space (section 25).
"""

from __future__ import annotations
import logging
import threading

logger = logging.getLogger("model_manager")
logging.basicConfig(level=logging.INFO)

_lock = threading.Lock()
_cache: dict[str, object] = {}

CLASSIFIER_MODEL_ID = "Luwayy/disaster_images_model"
CAPTION_MODEL_ID = "cnmoro/tiny-image-captioning"
LLM_MODEL_ID = "Qwen/Qwen2.5-0.5B-Instruct"


class ModelLoadError(Exception):
    pass


def _get_or_load(key: str, loader_fn):
    with _lock:
        if key in _cache:
            return _cache[key]
        try:
            logger.info(f"Loading model for '{key}' (first use, will be cached)...")
            obj = loader_fn()
            _cache[key] = obj
            logger.info(f"Model '{key}' loaded and cached.")
            return obj
        except Exception as e:
            logger.error(f"Failed to load model '{key}': {e}")
            raise ModelLoadError(f"Could not load model for {key}: {e}") from e


def get_classifier():
    """Image classification pipeline: disaster damage-type classification."""
    def _load():
        from transformers import pipeline
        return pipeline("image-classification", model=CLASSIFIER_MODEL_ID)
    return _get_or_load("classifier", _load)


def get_captioner():
    """Image captioning model: produces the 'evidence' text for a zone's photo."""
    def _load():
        from transformers import AutoTokenizer, AutoImageProcessor, VisionEncoderDecoderModel
        tokenizer = AutoTokenizer.from_pretrained(CAPTION_MODEL_ID)
        image_processor = AutoImageProcessor.from_pretrained(CAPTION_MODEL_ID)
        model = VisionEncoderDecoderModel.from_pretrained(CAPTION_MODEL_ID)
        model.eval()
        return {"model": model, "tokenizer": tokenizer, "image_processor": image_processor}
    return _get_or_load("captioner", _load)


def get_llm():
    """Small instruct LLM used for report parsing (JSON extraction) and explanation."""
    def _load():
        from transformers import AutoModelForCausalLM, AutoTokenizer
        tokenizer = AutoTokenizer.from_pretrained(LLM_MODEL_ID)
        model = AutoModelForCausalLM.from_pretrained(LLM_MODEL_ID, torch_dtype="auto")
        model.eval()
        return {"model": model, "tokenizer": tokenizer}
    return _get_or_load("llm", _load)


def release_all():
    """Frees cached models. Not called during normal demo operation, but available
    if the Space needs to recover memory between sessions."""
    with _lock:
        _cache.clear()
    logger.info("All cached models released.")


def preload_status() -> dict[str, bool]:
    """Reports which models are currently loaded, for the UI's system status strip."""
    return {
        "classifier": "classifier" in _cache,
        "captioner": "captioner" in _cache,
        "llm": "llm" in _cache,
    }