scripts/download_model.py

#!/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)