diff --git a/.env.example b/.env.example index 5da8bc8..eb1b3ab 100644 --- a/.env.example +++ b/.env.example @@ -116,7 +116,12 @@ SECRET_KEY=change-this-to-a-long-random-string-in-production CORS_ORIGINS=http://localhost:5173,http://localhost:3000,http://localhost:3080,http://localhost # Database path -DATABASE_URL=sqlite:///./data/photoedit.db +DATABASE_URL=sqlite:///./data/ai_photo_edit.db + +# Auto-download SAM model on startup (true/false) +# When true (default): Downloads SAM model (~375MB) on first startup for offline Smart Select +# When false: Skips download, Smart Select uses Replicate API (requires REPLICATE_API_KEY) +AUTO_DOWNLOAD_SAM=true # Allow users to select model per-edit ALLOW_MODEL_OVERRIDE=true diff --git a/backend/entrypoint.sh b/backend/entrypoint.sh index 362ca44..b44ca4d 100644 --- a/backend/entrypoint.sh +++ b/backend/entrypoint.sh @@ -3,9 +3,10 @@ # AI Photo Edit - Container Startup Script # ============================================================================= # This script runs when the container starts. It: -# 1. Downloads sample eye images if the catalog is empty -# 2. Ensures all directories exist -# 3. Starts the FastAPI server +# 1. Initializes the database +# 2. Downloads SAM model automatically (can be disabled with AUTO_DOWNLOAD_SAM=false) +# 3. Downloads sample eye images if the catalog is empty +# 4. Starts the FastAPI server # ============================================================================= set -e @@ -18,18 +19,15 @@ echo "==========================================" mkdir -p /app/data/projects mkdir -p /app/data/patches mkdir -p /app/data/models +mkdir -p /app/data/patch_library -# Check if eye catalog needs to be populated -echo "Checking eye catalog..." -PATCHES_COUNT=$(find /app/data/patches -maxdepth 1 -type d | wc -l) - -if [ "$PATCHES_COUNT" -le 1 ]; then - echo "Eye catalog is empty. Downloading sample eyes..." - python /scripts/download_sample_eyes.py || echo "Warning: Could not download sample eyes (non-fatal)" -else - echo "Eye catalog has content, skipping download." -fi +# Initialize database FIRST (before eye import) +echo "" +echo "Initializing database..." +echo "------------------------------------------" +cd /app && python /scripts/init_database.py || echo "Warning: Database init failed (non-fatal)" +# Check and download SAM model automatically echo "" echo "Checking SAM model (Smart Select)..." echo "------------------------------------------" @@ -39,17 +37,42 @@ if [ -f "/app/data/models/sam_model.pth" ] || \ [ -f "/app/data/models/sam_vit_h_4b8939.pth" ]; then echo "✓ SAM model found - Smart Select will use local AI (free, offline)" else - echo "" - echo "⚠ SAM model not found" - echo "" - echo " Smart Select will use Replicate API (requires REPLICATE_API_KEY)" - echo "" - echo " To enable FREE offline Smart Select, run:" - echo " docker exec -it ai-photo-edit-backend python /scripts/download_sam_model.py" - echo "" - echo " Model sizes: vit_b (375MB), vit_l (1.2GB), vit_h (2.5GB)" - echo " The model persists across container rebuilds." - echo "" + # Auto-download SAM unless explicitly disabled + AUTO_DOWNLOAD_SAM="${AUTO_DOWNLOAD_SAM:-true}" + if [ "$AUTO_DOWNLOAD_SAM" = "true" ]; then + echo "SAM model not found. Downloading automatically..." + echo "(This is a one-time ~375MB download that persists across rebuilds)" + echo "" + python /scripts/download_sam_model.py vit_b || { + echo "" + echo "⚠ SAM download failed (non-fatal)" + echo " Smart Select will fall back to Replicate API (requires REPLICATE_API_KEY)" + echo " To retry later: docker exec -it ai-photo-edit-backend python /scripts/download_sam_model.py" + } + else + echo "" + echo "⚠ SAM model not found (AUTO_DOWNLOAD_SAM=false)" + echo "" + echo " Smart Select will use Replicate API (requires REPLICATE_API_KEY)" + echo "" + echo " To enable FREE offline Smart Select, run:" + echo " docker exec -it ai-photo-edit-backend python /scripts/download_sam_model.py" + echo "" + fi +fi + +# Check if eye catalog needs to be populated +echo "" +echo "Checking eye catalog..." +echo "------------------------------------------" +PATCHES_COUNT=$(find /app/data/patches -maxdepth 1 -type d 2>/dev/null | wc -l) +DB_PATCHES_COUNT=$(sqlite3 /app/data/ai_photo_edit.db "SELECT COUNT(*) FROM patches;" 2>/dev/null || echo "0") + +if [ "$DB_PATCHES_COUNT" = "0" ] || [ "$PATCHES_COUNT" -le 1 ]; then + echo "Eye catalog is empty. Downloading sample eyes..." + cd /app && python /scripts/download_sample_eyes.py || echo "Warning: Could not download sample eyes (non-fatal)" +else + echo "✓ Eye catalog has $DB_PATCHES_COUNT patches" fi echo "" diff --git a/docker-compose.yml b/docker-compose.yml index ee51300..25c9539 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -20,6 +20,7 @@ services: - STABILITY_API_KEY=${STABILITY_API_KEY:-} - REPLICATE_API_KEY=${REPLICATE_API_KEY:-} - CORS_ORIGINS=http://localhost:5173,http://localhost:3000,http://localhost + - AUTO_DOWNLOAD_SAM=${AUTO_DOWNLOAD_SAM:-true} restart: unless-stopped networks: - ai-photo-edit-network diff --git a/scripts/download_sample_eyes.py b/scripts/download_sample_eyes.py index 6dacbf5..953aa42 100755 --- a/scripts/download_sample_eyes.py +++ b/scripts/download_sample_eyes.py @@ -106,17 +106,34 @@ def import_eye_to_database(db_path: Path, eye_data: dict, patch_path: str, thumb return patch_id def main(): - # Determine paths - script_dir = Path(__file__).parent - project_root = script_dir.parent - data_dir = project_root / 'data' - db_path = data_dir / 'photoedit.db' + # Determine paths - handle both Docker and local environments + # In Docker: script is at /scripts/, data is at /app/data/ + # Locally: script is at ./scripts/, data is at ./data/ + docker_data_dir = Path('/app/data') + local_data_dir = Path(__file__).parent.parent / 'data' + + if docker_data_dir.exists(): + data_dir = docker_data_dir + else: + data_dir = local_data_dir + + db_path = data_dir / 'ai_photo_edit.db' # Check if database exists if not db_path.exists(): print(f"Database not found at {db_path}") - print("Please start the backend first to initialize the database.") - sys.exit(1) + print("Attempting to initialize database...") + # Try to import and initialize database + try: + sys.path.insert(0, str(Path('/app'))) + sys.path.insert(0, str(Path(__file__).parent.parent / 'backend')) + from app.database import init_db + init_db() + print("Database initialized successfully.") + except Exception as e: + print(f"Could not initialize database: {e}") + print("Please start the backend first to initialize the database.") + sys.exit(1) print(f"Using database: {db_path}") print(f"Data directory: {data_dir}") diff --git a/scripts/init_database.py b/scripts/init_database.py new file mode 100644 index 0000000..aa99aef --- /dev/null +++ b/scripts/init_database.py @@ -0,0 +1,45 @@ +#!/usr/bin/env python3 +""" +Initialize the database before other startup scripts run. + +This ensures the database exists and has all required tables +before download_sample_eyes.py tries to use it. +""" + +import os +import sys +from pathlib import Path + +# Add backend to path - handle both Docker and local environments +# In Docker: backend is at /app/ +# Locally: backend is at ./backend/ +if Path('/app').exists(): + sys.path.insert(0, '/app') +else: + sys.path.insert(0, str(Path(__file__).parent.parent / 'backend')) + +def main(): + # Import after path setup + from app.database import engine, Base, init_db + from app.models import project, user, patch + + print("Initializing database...") + + # Create all tables + init_db() + + # Verify database was created - check both Docker and local paths + docker_db_path = Path('/app/data/ai_photo_edit.db') + local_db_path = Path('./data/ai_photo_edit.db') + + if docker_db_path.exists(): + print(f"✓ Database initialized at: {docker_db_path}") + elif local_db_path.exists(): + print(f"✓ Database initialized at: {local_db_path}") + else: + print("⚠ Database file not found at expected locations, but tables may still be created") + + print("Database initialization complete.") + +if __name__ == '__main__': + main()