prefetch_models.py

"""
Pre-download and cache both AI models INTO the Docker image at build time.

Why: without this, the first real user request after a deploy has to download
~350MB (SAM vit_b) + ~170MB (Mask2Former swin-tiny) over the network AND load
them into memory, all within a single HTTP request — which can exceed the
platform's request timeout and return a 503/timeout before the app even logs
the request.

Running this during `docker build` means both models are already on disk
when the container starts, so the first request only has to load them into
RAM (fast), not download them (slow).
"""
import logging
from pathlib import Path

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

BASE_DIR = Path(__file__).resolve().parent

# ── SAM (vit_b) ──────────────────────────────────────────────────────
SAM_DIR = BASE_DIR / "sam"
SAM_DIR.mkdir(parents=True, exist_ok=True)
SAM_CHECKPOINT = SAM_DIR / "sam_vit_b_01ec64.pth"

if not SAM_CHECKPOINT.exists():
    logger.info("Downloading SAM vit_b checkpoint...")
    from huggingface_hub import hf_hub_download
    hf_hub_download(
        repo_id="GraydientPlatformAPI/sams",
        filename="sam_vit_b_01ec64.pth",
        local_dir=str(SAM_DIR),
    )
    logger.info("SAM vit_b checkpoint cached.")
else:
    logger.info("SAM vit_b checkpoint already present.")

# ── Mask2Former (swin-tiny) ──────────────────────────────────────────
MASK2FORMER_CACHE = BASE_DIR / "models" / "mask2former"
MASK2FORMER_CACHE.mkdir(parents=True, exist_ok=True)

if not (MASK2FORMER_CACHE / "config.json").exists():
    logger.info("Downloading Mask2Former (swin-tiny)...")
    from transformers import Mask2FormerForUniversalSegmentation, AutoImageProcessor

    model_id = "facebook/mask2former-swin-tiny-ade-semantic"
    processor = AutoImageProcessor.from_pretrained(model_id)
    model = Mask2FormerForUniversalSegmentation.from_pretrained(model_id)

    processor.save_pretrained(str(MASK2FORMER_CACHE))
    model.save_pretrained(str(MASK2FORMER_CACHE))
    logger.info("Mask2Former cached.")
else:
    logger.info("Mask2Former already cached.")

logger.info("Prefetch complete — both models are baked into the image.")