Improve Hugging Face model download process: add error handling for token login and model download, log success and failure messages, and ensure model is downloaded to the correct directory.

app/main.py

@@ -127,19 +127,35 @@ async def lifespan(app: FastAPI):
         try:
             from huggingface_hub import hf_hub_download, login
             # Check for token in environment
-            hf_token = os.environ.get("HF_TOKEN")
+            hf_token = os.environ.get("HF_TOKEN", "").strip()
             if hf_token:
                 logger.info("Logging in to Hugging Face using provided token")
-                login(token=hf_token)
-            logger.info("Downloading CSM-1B model from Hugging Face...")
-            download_start = time.time()
-            model_path = hf_hub_download(
-                repo_id="sesame/csm-1b", 
-                filename="ckpt.pt", 
-                local_dir=APP_DIR
-            )
-            download_time = time.time() - download_start
-            logger.info(f"Model downloaded to {model_path} in {download_time:.2f} seconds")
+                try:
+                    login(token=hf_token)
+                    logger.info("Successfully logged in to Hugging Face")
+                except Exception as e:
+                    logger.error(f"Error logging in to Hugging Face: {e}")
+                    logger.error("Will attempt to download model without authentication")
+            else:
+                logger.warning("No Hugging Face token provided. Some models may not be accessible")
+
+            # Attempt to download the model
+            try:
+                logger.info("Downloading CSM-1B model from Hugging Face...")
+                download_start = time.time()
+                model_path = hf_hub_download(
+                    repo_id="sesame/csm-1b", 
+                    filename="ckpt.pt",
+                    local_dir=os.path.join(APP_DIR, "models"),
+                    token=hf_token if hf_token else None
+                )
+                download_time = time.time() - download_start
+                logger.info(f"Model downloaded to {model_path} in {download_time:.2f} seconds")
+            except Exception as e:
+                error_stack = traceback.format_exc()
+                logger.error(f"Error downloading model: {str(e)}\n{error_stack}")
+                logger.error("Please ensure you have a valid Hugging Face token with access to the model")
+                logger.error("Starting without model - API will return 503 Service Unavailable")
         except Exception as e:
             error_stack = traceback.format_exc()
             logger.error(f"Error downloading model: {str(e)}\n{error_stack}")

scripts/download_model.py

@@ -0,0 +1,60 @@
+#!/usr/bin/env python3
+"""
+Script to download the CSM-1B model from Hugging Face.
+"""
+import os
+import sys
+import logging
+from huggingface_hub import hf_hub_download, login
+from pathlib import Path
+
+# Configure logging
+logging.basicConfig(
+    level=logging.INFO,
+    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
+)
+logger = logging.getLogger(__name__)
+
+def download_model():
+    """Download the CSM-1B model."""
+    try:
+        # Set up paths
+        app_dir = Path("/app")
+        models_dir = app_dir / "models"
+        models_dir.mkdir(parents=True, exist_ok=True)
+        
+        # Get token from environment
+        hf_token = os.environ.get("HF_TOKEN", "").strip()
+        
+        # Try to login if token is provided
+        if hf_token:
+            logger.info("Logging in to Hugging Face...")
+            try:
+                login(token=hf_token)
+                logger.info("Successfully logged in to Hugging Face")
+            except Exception as e:
+                logger.error(f"Error logging in to Hugging Face: {e}")
+                logger.error("Will attempt to download without authentication")
+        else:
+            logger.warning("No Hugging Face token provided. Some models may not be accessible")
+        
+        # Download the model
+        logger.info("Downloading CSM-1B model...")
+        model_path = hf_hub_download(
+            repo_id="sesame/csm-1b",
+            filename="ckpt.pt",
+            local_dir=str(models_dir),
+            token=hf_token if hf_token else None
+        )
+        
+        logger.info(f"Model downloaded successfully to: {model_path}")
+        logger.info(f"File size: {os.path.getsize(model_path) / (1024*1024):.2f} MB")
+        return True
+        
+    except Exception as e:
+        logger.error(f"Error downloading model: {e}")
+        return False
+
+if __name__ == "__main__":
+    success = download_model()
+    sys.exit(0 if success else 1)