Added the feature to check and download models in Space from the actual model repository.

app.py

@@ -3,10 +3,12 @@ import gradio as gr
 import torch
 from TeLVE.imagine import ImageCaptioningModel, load_model, generate_caption
 from transformers import BertTokenizerFast
+from huggingface_hub import hf_hub_download, list_repo_files
 
 # Constants
 MODELS_DIR = "./TeLVE/models"
 TOKENIZER_PATH = "./TeLVE/tokenizer"
+HF_REPO_ID = "outsu/TeLVE"
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 def list_available_models():
@@ -15,6 +17,35 @@ def list_available_models():
         return []
     return [f for f in os.listdir(MODELS_DIR) if f.endswith('.pth')]
 
+def get_hf_model_list():
+    """Get list of model files from HuggingFace repository"""
+    try:
+        files = list_repo_files(HF_REPO_ID)
+        return [f.split('/')[-1] for f in files if f.startswith('models/') and f.endswith('.pth')]
+    except Exception as e:
+        print(f"Error fetching models from HuggingFace: {str(e)}")
+        return []
+
+def download_missing_models():
+    """Download missing models from HuggingFace"""
+    if not os.path.exists(MODELS_DIR):
+        os.makedirs(MODELS_DIR)
+    
+    local_models = set(list_available_models())
+    hf_models = set(get_hf_model_list())
+    
+    for model in hf_models - local_models:
+        try:
+            print(f"Downloading missing model: {model}")
+            hf_hub_download(
+                repo_id=HF_REPO_ID,
+                filename=f"models/{model}",
+                local_dir=os.path.dirname(MODELS_DIR),
+                local_dir_use_symlinks=False
+            )
+        except Exception as e:
+            print(f"Error downloading {model}: {str(e)}")
+
 def generate_description(image, model_name):
     """Generate image caption using selected model"""
     try:
@@ -65,5 +96,7 @@ def create_interface():
     return interface
 
 if __name__ == "__main__":
+    print("Checking for missing models...")
+    download_missing_models()
     demo = create_interface()
     demo.launch(share=True, server_name="0.0.0.0")

requirements.txt

@@ -1,8 +1,8 @@
-
 torch>=1.9.0
 torchvision>=0.10.0
 transformers>=4.11.0
 gradio>=3.0.0
 pandas>=1.3.0
 Pillow>=8.0.0
-tqdm>=4.62.0
+tqdm>=4.62.0
+huggingface-hub>=0.16.4