Spaces:
Running
Running
File size: 4,602 Bytes
32d4a86 c8cbd99 32d4a86 c8cbd99 32d4a86 c8cbd99 32d4a86 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 |
"""
Download EfficientNet model from cloud storage if not present.
This script runs at application startup to download the model if needed.
"""
import os
import sys
import logging
from pathlib import Path
import requests
logger = logging.getLogger(__name__)
def download_efficientnet_model():
"""
Download EfficientNet optimized model if it doesn't exist.
Supports two methods:
1. Hugging Face Hub (set HF_MODEL_REPO environment variable)
2. Direct URL download (set EFFICIENTNET_MODEL_URL environment variable)
"""
# Get base directory
base_dir = Path(__file__).resolve().parent.parent
models_dir = base_dir / "models" / "optimized_models"
models_dir.mkdir(parents=True, exist_ok=True)
model_path = models_dir / "efficientnet_efficient_best_model_quantized.pth"
# Check if model already exists
if model_path.exists():
size_mb = model_path.stat().st_size / (1024 * 1024)
logger.info(f"EfficientNet model already exists ({size_mb:.1f}MB)")
return True
# Try Hugging Face Hub first
hf_repo = os.environ.get("HF_MODEL_REPO")
if hf_repo:
try:
from huggingface_hub import hf_hub_download
logger.info(f"Downloading model from Hugging Face Hub: {hf_repo}")
downloaded_path = hf_hub_download(
repo_id=hf_repo,
filename="efficientnet_efficient_best_model_quantized.pth",
cache_dir=str(models_dir),
local_dir=str(models_dir),
local_dir_use_symlinks=False
)
# Copy to expected location if needed (HF Hub creates nested structure)
downloaded_path = Path(downloaded_path)
if downloaded_path != model_path:
import shutil
logger.info(f"Copying model from {downloaded_path} to {model_path}")
shutil.copy2(downloaded_path, model_path)
logger.info(f"Model copied to expected location: {model_path}")
size_mb = model_path.stat().st_size / (1024 * 1024)
logger.info(f"Model downloaded from HF Hub successfully ({size_mb:.1f}MB) at {model_path}")
return True
except ImportError:
logger.warning("huggingface_hub not installed. Install with: pip install huggingface_hub")
except Exception as e:
logger.warning(f"Failed to download from HF Hub: {e}. Trying direct URL...")
# Fallback to direct URL download
model_url = os.environ.get("EFFICIENTNET_MODEL_URL")
if not model_url:
logger.warning("Neither HF_MODEL_REPO nor EFFICIENTNET_MODEL_URL is set.")
logger.warning("Model will not be downloaded. Set one of these environment variables.")
return False
try:
logger.info(f"Downloading EfficientNet model from {model_url}...")
logger.info("This may take a few minutes (model is ~245MB)...")
# Download with progress
response = requests.get(model_url, stream=True, timeout=300)
response.raise_for_status()
total_size = int(response.headers.get('content-length', 0))
downloaded = 0
with open(model_path, 'wb') as f:
for chunk in response.iter_content(chunk_size=8192):
if chunk:
f.write(chunk)
downloaded += len(chunk)
if total_size > 0:
percent = (downloaded / total_size) * 100
if downloaded % (10 * 1024 * 1024) == 0: # Log every 10MB
logger.info(f"Downloaded {downloaded / (1024 * 1024):.1f}MB / {total_size / (1024 * 1024):.1f}MB ({percent:.1f}%)")
size_mb = model_path.stat().st_size / (1024 * 1024)
logger.info(f"EfficientNet model downloaded successfully ({size_mb:.1f}MB)")
return True
except requests.exceptions.RequestException as e:
logger.error(f"Failed to download model: {e}")
# Clean up partial download
if model_path.exists():
model_path.unlink()
return False
except Exception as e:
logger.error(f"Error downloading model: {e}", exc_info=True)
# Clean up partial download
if model_path.exists():
model_path.unlink()
return False
if __name__ == "__main__":
# Configure logging
logging.basicConfig(
level=logging.INFO,
format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
download_efficientnet_model()
|