Update app.py

app.py

@@ -7,6 +7,8 @@ import json
 import os
 from datetime import datetime
 import uuid
+from huggingface_hub import snapshot_download
+import shutil
 
 # Set up logging
 logging.basicConfig(level=logging.INFO)
@@ -15,14 +17,27 @@ logger = logging.getLogger(__name__)
 # Define output JSON file path
 OUTPUT_JSON_PATH = "captions.json"
 
-# Load the model and processor
+# Clear Hugging Face cache and download model
 MODEL_PATH = "fancyfeast/llama-joycaption-beta-one-hf-llava"
 try:
+    # Clear cache to avoid corrupted files
+    cache_dir = os.path.expanduser("~/.cache/huggingface/hub")
+    model_cache = os.path.join(cache_dir, f"models--{MODEL_PATH.replace('/', '--')}")
+    if os.path.exists(model_cache):
+        shutil.rmtree(model_cache)
+        logger.info(f"Cleared cache for {MODEL_PATH}")
+
+    # Pre-download model to ensure integrity
+    snapshot_download(repo_id=MODEL_PATH)
+    logger.info(f"Downloaded model {MODEL_PATH}")
+
+    # Load processor and model
     processor = AutoProcessor.from_pretrained(MODEL_PATH)
     model = LlavaForConditionalGeneration.from_pretrained(
         MODEL_PATH,
         torch_dtype=torch.float32,  # CPU-compatible dtype
-        low_cpu_mem_usage=True     # Minimize memory usage
+        low_cpu_mem_usage=True,    # Minimize memory usage
+        use_safetensors=True       # Force safetensors
     ).to("cpu")
     model.eval()
     logger.info("Model and processor loaded successfully.")