feat: load model from HuggingFace Hub

app.py

@@ -5,7 +5,7 @@ from torchvision import transforms
 from PIL import Image
 import numpy as np
 import os
-from urllib.request import urlretrieve
+from huggingface_hub import hf_hub_download
 
 # --- Model Definition ---
 
@@ -241,7 +241,7 @@ device = None
 
 
 def load_model():
-    """Loads the model and caches it in a global variable."""
+    """Loads the model from HuggingFace Hub and caches it in a global variable."""
     global model, device
     if model is not None:
         return model, device
@@ -249,19 +249,19 @@ def load_model():
     try:
         device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-        MODEL_PATH = "best_model.pth"
+        # Download model from HuggingFace Hub
+        model_path = hf_hub_download(
+            repo_id="divitmittal/HybridTransformer-MFIF",
+            filename="best_model.pth",
+            cache_dir="./model_cache"
+        )
 
         model_instance = FocalCrossViTHybrid(img_size=224).to(device)
 
-        if not os.path.exists(MODEL_PATH):
-            raise FileNotFoundError(
-                "Model checkpoint 'best_model.pth' not found. Please ensure it is available in the Space."
-            )
-
-        checkpoint = torch.load(MODEL_PATH, map_location=device)
+        checkpoint = torch.load(model_path, map_location=device)
 
         state_dict = checkpoint.get("model_state_dict", checkpoint)
-        if any(key.startswith("module.")):
+        if any(key.startswith("module.") for key in state_dict.keys()):
             state_dict = {k.replace("module.", ""): v for k, v in state_dict.items()}
 
         model_instance.load_state_dict(state_dict)
@@ -271,7 +271,7 @@ def load_model():
         return model, device
     except Exception as e:
         # Catch any exception during loading and show it in the UI
-        raise gr.Error(f"Failed to load the model: {e}")
+        raise gr.Error(f"Failed to load the model from HuggingFace Hub: {e}")
 
 
 # Image processing functions

requirements.txt

@@ -3,3 +3,4 @@ torchvision
 gradio
 numpy
 Pillow
+huggingface_hub