Update main.py

main.py

@@ -6,6 +6,8 @@ import torch
 import pickle
 import os
 import uvicorn
+from huggingface_hub import hf_hub_download
+
 # Import from model.py
 from model import (
     Vocabulary,
@@ -34,12 +36,16 @@ app.add_middleware(
 DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 # -------------------------
-# Paths (relative to main.py)
+# Paths
 # -------------------------
 BASE_DIR = os.path.dirname(os.path.abspath(__file__))
 
 VOCAB_PATH = os.path.join(BASE_DIR, "vocab.pkl")
-CHECKPOINT_PATH = os.path.join(BASE_DIR, "best_checkpoint.pth")
+
+CHECKPOINT_PATH = hf_hub_download(
+    repo_id="VIKRAM989/image-label",
+    filename="best_checkpoint.pth"
+)
 
 # -------------------------
 # Load Vocabulary
@@ -67,6 +73,7 @@ model = ImageCaptioningModel(encoder, decoder).to(DEVICE)
 # Load Weights
 # -------------------------
 checkpoint = torch.load(CHECKPOINT_PATH, map_location=DEVICE)
+
 model.load_state_dict(checkpoint["model_state_dict"])
 
 model.eval()
@@ -94,10 +101,8 @@ async def caption_image(file: UploadFile = File(...)):
 
     caption = generate_caption(model, image, vocab)
 
-    return {
-        "caption": caption
-    }
+    return {"caption": caption}
 
-if __name__ == "__main__":
 
+if __name__ == "__main__":
     uvicorn.run("main:app", host="0.0.0.0", port=7860)