main.py

from fastapi import FastAPI
from fastapi.responses import HTMLResponse
import uvicorn
from vit_captioning.generate import CaptionGenerator

app = FastAPI()
caption_generator = None  # Lazy-load placeholder

@app.on_event("startup")
def startup_event():
    global caption_generator
    if caption_generator is None:
        print("Loading CaptionGenerator...")
        caption_generator = CaptionGenerator(
            model_type="CLIPEncoder",
            checkpoint_path="./vit_captioning/artifacts/CLIPEncoder_40epochs_unfreeze12.pth",
            quantized=False
        )

@app.get("/", response_class=HTMLResponse)
def root():
    return "<h3>✅ Hugging Face Space is alive</h3>"

@app.get("/health")
def health_check():
    return {"status": "ok"}

# Example endpoint to trigger model
@app.get("/caption")
def caption():
    if caption_generator is None:
        return {"error": "Model not loaded"}
    return {"result": "dummy caption"}  # Replace with real logic

# if __name__ == "__main__":
#     uvicorn.run(app, host="0.0.0.0", port=8000)