import os
import torch
from transformers import (
    DetrImageProcessor, 
    DetrForObjectDetection,
    ViTImageProcessor,
    ViTForImageClassification,
    CLIPProcessor,
    CLIPModel,
    AutoTokenizer,
    AutoModel,
    AutoModelForQuestionAnswering,
    AutoModelForSeq2SeqLM,
    BartForConditionalGeneration
)

# Set timeout
os.environ["HF_HUB_DOWNLOAD_TIMEOUT"] = "120"

def preload():
    print("🚀 Starting model pre-loading...")
    
    models = {
        "detection": ("facebook/detr-resnet-50", DetrForObjectDetection, DetrImageProcessor),
        "reid": ("google/vit-base-patch16-224", ViTForImageClassification, ViTImageProcessor),
        "clip": ("openai/clip-vit-base-patch32", CLIPModel, CLIPProcessor),
        "search": ("sentence-transformers/all-MiniLM-L6-v2", AutoModel, AutoTokenizer),
        "qa": ("deepset/roberta-base-squad2", AutoModelForQuestionAnswering, AutoTokenizer),
        "report": ("google/flan-t5-base", AutoModelForSeq2SeqLM, AutoTokenizer),
        "summarizer": ("facebook/bart-large-cnn", BartForConditionalGeneration, AutoTokenizer),
    }

    for name, (model_id, model_cls, proc_cls) in models.items():
        print(f"📦 Pre-loading {name}: {model_id}...")
        try:
            proc_cls.from_pretrained(model_id)
            model_cls.from_pretrained(model_id)
            print(f"✅ {name} loaded.")
        except Exception as e:
            print(f"❌ Failed to load {name}: {e}")

if __name__ == "__main__":
    preload()