| """ |
| 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_DIR = BASE_DIR / "sam" |
| SAM_DIR.mkdir(parents=True, exist_ok=True) |
| SAM_CHECKPOINT = SAM_DIR / "sam_vit_h_4b8939.pth" |
|
|
| if not SAM_CHECKPOINT.exists(): |
| logger.info("Downloading SAM vit_h checkpoint (2.5GB, this will take a while)...") |
| from huggingface_hub import hf_hub_download |
| hf_hub_download( |
| repo_id="HCMUE-Research/SAM-vit-h", |
| filename="sam_vit_h_4b8939.pth", |
| local_dir=str(SAM_DIR), |
| ) |
| logger.info("SAM vit_h checkpoint cached.") |
| else: |
| logger.info("SAM vit_h checkpoint already present.") |
|
|
| |
| 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.") |
|
|
