Upload app.py

app.py

@@ -1,4 +1,5 @@
 import os
+import json
 import gradio as gr
 import torch
 from TeLVE.imagine import ImageCaptioningModel, load_model, generate_caption
@@ -9,8 +10,22 @@ from huggingface_hub import hf_hub_download, list_repo_files
 MODELS_DIR = "./TeLVE/models"
 TOKENIZER_PATH = "./TeLVE/tokenizer"
 HF_REPO_ID = "outsu/TeLVE"
+MODEL_STATE_FILE = "./TeLVE/model_state.json"
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
+def load_model_state():
+    """Load the model state from JSON file"""
+    if os.path.exists(MODEL_STATE_FILE):
+        with open(MODEL_STATE_FILE, 'r') as f:
+            return json.load(f)
+    return {"downloaded_models": []}
+
+def save_model_state(state):
+    """Save the model state to JSON file"""
+    os.makedirs(os.path.dirname(MODEL_STATE_FILE), exist_ok=True)
+    with open(MODEL_STATE_FILE, 'w') as f:
+        json.dump(state, f)
+
 def list_available_models():
     """List all .pth models in the models directory"""
     if not os.path.exists(MODELS_DIR):
@@ -27,14 +42,19 @@ def get_hf_model_list():
         return []
 
 def download_missing_models():
-    """Download missing models from HuggingFace"""
+    """Download missing models from HuggingFace with state management"""
     if not os.path.exists(MODELS_DIR):
         os.makedirs(MODELS_DIR)
     
+    state = load_model_state()
+    downloaded_models = set(state["downloaded_models"])
     local_models = set(list_available_models())
     hf_models = set(get_hf_model_list())
     
-    for model in hf_models - local_models:
+    # Check for models that need downloading
+    models_to_download = (hf_models - local_models) - downloaded_models
+    
+    for model in models_to_download:
         try:
             print(f"Downloading missing model: {model}")
             hf_hub_download(
@@ -43,8 +63,21 @@ def download_missing_models():
                 local_dir=os.path.dirname(MODELS_DIR),
                 local_dir_use_symlinks=False
             )
+            downloaded_models.add(model)
         except Exception as e:
             print(f"Error downloading {model}: {str(e)}")
+            continue
+    
+    # Update state with newly downloaded models
+    state["downloaded_models"] = list(downloaded_models)
+    save_model_state(state)
+
+def verify_model_integrity():
+    """Verify that all models in state actually exist"""
+    state = load_model_state()
+    local_models = set(list_available_models())
+    state["downloaded_models"] = list(set(state["downloaded_models"]) & local_models)
+    save_model_state(state)
 
 def generate_description(image, model_name):
     """Generate image caption using selected model"""
@@ -96,6 +129,8 @@ def create_interface():
     return interface
 
 if __name__ == "__main__":
+    print("Verifying model integrity...")
+    verify_model_integrity()
     print("Checking for missing models...")
     download_missing_models()
     demo = create_interface()