Merge pull request #12 from outis1one/claude/add-eye-detection-feature-69XOl
Claude/add eye detection feature 69 x ol
This commit is contained in:
outis1one
+142
-22
@@ -1,28 +1,148 @@
|
||||
# AI Provider Configuration
|
||||
# Options: openai, stability, replicate, mock
|
||||
# - openai: DALL-E 2 (low quality, not recommended)
|
||||
# - stability: Stability AI SDXL (good quality, ~$0.04/image)
|
||||
# - replicate: Multiple models (best value, ~$0.002-0.025/image)
|
||||
# - mock: No AI, returns original (for testing)
|
||||
AI_PROVIDER=mock
|
||||
# =============================================================================
|
||||
# AI Photo Edit - Environment Configuration
|
||||
# =============================================================================
|
||||
#
|
||||
# SETUP INSTRUCTIONS:
|
||||
# 1. Copy this file to .env: cp .env.example .env
|
||||
# 2. Get API key from Replicate (see below)
|
||||
# 3. Paste your key in the REPLICATE_API_KEY line
|
||||
# 4. Rebuild: docker-compose up -d --build
|
||||
#
|
||||
# =============================================================================
|
||||
|
||||
# API Keys
|
||||
OPENAI_API_KEY=
|
||||
STABILITY_API_KEY=
|
||||
REPLICATE_API_KEY=
|
||||
|
||||
# Model Selection (optional, provider-specific)
|
||||
# Stability AI models: sdxl (default), sd15, sd21
|
||||
STABILITY_MODEL=sdxl
|
||||
# =============================================================================
|
||||
# STEP 1: Choose AI Provider
|
||||
# =============================================================================
|
||||
# Options: mock, openai, stability, replicate
|
||||
#
|
||||
# mock = Free, but returns original image unchanged (for testing UI)
|
||||
# openai = DALL-E 2 inpainting (~$0.02/image) - lower quality
|
||||
# stability = Stability AI SDXL (~$0.01/image) - good quality
|
||||
# replicate = Multiple models (~$0.002-0.03/image) - RECOMMENDED
|
||||
#
|
||||
# RECOMMENDED: Use "replicate" for best quality and model variety
|
||||
# =============================================================================
|
||||
|
||||
# Replicate models: sdxl-inpaint (default), lama, realistic-vision
|
||||
# - sdxl-inpaint: Best general purpose (~$0.025/image)
|
||||
# - lama: Best for object removal (~$0.002/image)
|
||||
# - realistic-vision: Best for humans/faces/hands (~$0.020/image)
|
||||
REPLICATE_MODEL=sdxl-inpaint
|
||||
AI_PROVIDER=replicate
|
||||
|
||||
# Allow per-edit model override (true/false)
|
||||
|
||||
# =============================================================================
|
||||
# STEP 2: Get Your API Key
|
||||
# =============================================================================
|
||||
#
|
||||
# ╔═══════════════════════════════════════════════════════════════════════════╗
|
||||
# ║ REPLICATE (RECOMMENDED) ║
|
||||
# ╠═══════════════════════════════════════════════════════════════════════════╣
|
||||
# ║ ║
|
||||
# ║ 1. Go to: https://replicate.com ║
|
||||
# ║ 2. Click "Sign in" (use GitHub, Google, or email) ║
|
||||
# ║ 3. Go to: https://replicate.com/account/api-tokens ║
|
||||
# ║ 4. Click "Create token" ║
|
||||
# ║ 5. Copy the token (starts with "r8_") ║
|
||||
# ║ 6. Paste it below after REPLICATE_API_KEY= ║
|
||||
# ║ ║
|
||||
# ║ FREE TIER: New accounts get some free credits to try models! ║
|
||||
# ║ PRICING: ~$0.002-0.03 per image depending on model ║
|
||||
# ║ ║
|
||||
# ╚═══════════════════════════════════════════════════════════════════════════╝
|
||||
|
||||
REPLICATE_API_KEY=r8_PASTE_YOUR_KEY_HERE
|
||||
|
||||
# ───────────────────────────────────────────────────────────────────────────
|
||||
# OPENAI (Alternative - not recommended, lower quality)
|
||||
# Get key at: https://platform.openai.com/api-keys
|
||||
# ───────────────────────────────────────────────────────────────────────────
|
||||
#OPENAI_API_KEY=sk-xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx
|
||||
|
||||
# ───────────────────────────────────────────────────────────────────────────
|
||||
# STABILITY AI (Alternative)
|
||||
# Get key at: https://platform.stability.ai/account/keys
|
||||
# ───────────────────────────────────────────────────────────────────────────
|
||||
#STABILITY_API_KEY=sk-xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# STEP 3: Model Selection (OPTIONAL - for advanced users)
|
||||
# =============================================================================
|
||||
#
|
||||
# By default, the system AUTO-SELECTS the best model based on your prompt:
|
||||
# - Prompt contains "remove/erase/delete" → Uses LaMa (fast removal)
|
||||
# - Prompt contains "face/hands/person" → Uses Realistic Vision
|
||||
# - Everything else → Uses SDXL Inpaint
|
||||
#
|
||||
# To FORCE a specific model, uncomment ONE line below:
|
||||
# ───────────────────────────────────────────────────────────────────────────
|
||||
|
||||
# REPLICATE_MODEL=sdxl-inpaint # General purpose, good quality (~$0.01)
|
||||
# REPLICATE_MODEL=lama # Object removal ONLY (~$0.002, fastest)
|
||||
# REPLICATE_MODEL=realistic-vision # Faces, hands, skin (~$0.02)
|
||||
|
||||
# ───────────────────────────────────────────────────────────────────────────
|
||||
# IMPORTANT: About Flux and other text-to-image models
|
||||
# ───────────────────────────────────────────────────────────────────────────
|
||||
#
|
||||
# Models like "black-forest-labs/flux-kontext-pro" are TEXT-TO-IMAGE models.
|
||||
# They generate NEW images from text, they DON'T edit existing images.
|
||||
#
|
||||
# For EDITING (inpainting), you need models that accept:
|
||||
# - An existing image
|
||||
# - A mask showing what to change
|
||||
# - A prompt describing the change
|
||||
#
|
||||
# WORKS for editing: DOESN'T work for editing:
|
||||
# ✓ sdxl-inpaint ✗ flux-kontext-pro (text-to-image)
|
||||
# ✓ lama ✗ flux-dev (text-to-image)
|
||||
# ✓ realistic-vision ✗ ideogram (text-to-image)
|
||||
#
|
||||
# ───────────────────────────────────────────────────────────────────────────
|
||||
|
||||
# Stability AI model selection (if using AI_PROVIDER=stability)
|
||||
#STABILITY_MODEL=sdxl # Options: sdxl, sd15, sd21
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# SECURITY (Change this in production!)
|
||||
# =============================================================================
|
||||
|
||||
SECRET_KEY=change-this-to-a-long-random-string-in-production
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# ADVANCED SETTINGS (Usually don't need to change)
|
||||
# =============================================================================
|
||||
|
||||
# CORS origins (comma-separated)
|
||||
CORS_ORIGINS=http://localhost:5173,http://localhost:3000,http://localhost:3080,http://localhost
|
||||
|
||||
# Database path
|
||||
DATABASE_URL=sqlite:///./data/photoedit.db
|
||||
|
||||
# Allow users to select model per-edit
|
||||
ALLOW_MODEL_OVERRIDE=true
|
||||
|
||||
# Secret key for JWT tokens (change in production)
|
||||
SECRET_KEY=change-this-secret-key-in-production
|
||||
|
||||
# =============================================================================
|
||||
# TROUBLESHOOTING
|
||||
# =============================================================================
|
||||
#
|
||||
# PROBLEM: "405 Method Not Allowed" errors
|
||||
# FIX: Rebuild container: docker-compose build --no-cache && docker-compose up -d
|
||||
#
|
||||
# PROBLEM: "REPLICATE_API_KEY not configured"
|
||||
# FIX: 1. Make sure .env file exists (not just .env.example)
|
||||
# 2. Make sure REPLICATE_API_KEY has your actual key
|
||||
# 3. Restart: docker-compose down && docker-compose up -d
|
||||
#
|
||||
# PROBLEM: Edits don't change the image
|
||||
# FIX: Check AI_PROVIDER isn't set to "mock"
|
||||
#
|
||||
# PROBLEM: "rembg not installed"
|
||||
# FIX: Rebuild: docker-compose build --no-cache backend
|
||||
#
|
||||
# PROBLEM: Smart Select uses flood-fill instead of AI
|
||||
# FIX: Smart Select needs REPLICATE_API_KEY for SAM model
|
||||
#
|
||||
# CHECK LOGS: docker-compose logs -f backend
|
||||
#
|
||||
# =============================================================================
|
||||
|
||||
+6
-2
@@ -25,11 +25,15 @@ RUN python -c "from rembg import remove; print('rembg model downloaded')" || tru
|
||||
# Copy application
|
||||
COPY . .
|
||||
|
||||
# Copy entrypoint script and make it executable
|
||||
COPY entrypoint.sh /entrypoint.sh
|
||||
RUN chmod +x /entrypoint.sh
|
||||
|
||||
# Create data directories
|
||||
RUN mkdir -p /app/data/projects /app/data/patches /app/data/models
|
||||
|
||||
# Expose port
|
||||
EXPOSE 8000
|
||||
|
||||
# Run the application
|
||||
CMD ["uvicorn", "app.main:app", "--host", "0.0.0.0", "--port", "8000"]
|
||||
# Use entrypoint script (auto-populates eyes on first run, then starts server)
|
||||
ENTRYPOINT ["/entrypoint.sh"]
|
||||
|
||||
+117
-19
@@ -172,31 +172,138 @@ async def smart_select(
|
||||
)
|
||||
|
||||
|
||||
# Global SAM model cache (loaded once, reused)
|
||||
_sam_model = None
|
||||
_sam_predictor = None
|
||||
|
||||
|
||||
def _get_sam_model():
|
||||
"""Load SAM model from local file (cached after first load)"""
|
||||
global _sam_model, _sam_predictor
|
||||
|
||||
if _sam_predictor is not None:
|
||||
return _sam_predictor
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
# Check for SAM model in models directory
|
||||
models_dir = Path('/app/data/models')
|
||||
model_path = models_dir / 'sam_model.pth'
|
||||
|
||||
# Also check for specific model files
|
||||
if not model_path.exists():
|
||||
for filename in ['sam_vit_b_01ec64.pth', 'sam_vit_l_0b3195.pth', 'sam_vit_h_4b8939.pth']:
|
||||
alt_path = models_dir / filename
|
||||
if alt_path.exists():
|
||||
model_path = alt_path
|
||||
break
|
||||
|
||||
if not model_path.exists():
|
||||
raise FileNotFoundError(
|
||||
f"SAM model not found. Download it with:\n"
|
||||
f" docker exec -it ai-photo-edit-backend python /scripts/download_sam_model.py"
|
||||
)
|
||||
|
||||
# Determine model type from filename
|
||||
model_type = 'vit_b' # default
|
||||
if 'vit_l' in model_path.name:
|
||||
model_type = 'vit_l'
|
||||
elif 'vit_h' in model_path.name:
|
||||
model_type = 'vit_h'
|
||||
|
||||
print(f"Loading SAM model: {model_path} (type: {model_type})")
|
||||
|
||||
import torch
|
||||
from segment_anything import sam_model_registry, SamPredictor
|
||||
|
||||
# Use CPU by default (works everywhere), GPU if available
|
||||
device = 'cuda' if torch.cuda.is_available() else 'cpu'
|
||||
|
||||
_sam_model = sam_model_registry[model_type](checkpoint=str(model_path))
|
||||
_sam_model.to(device)
|
||||
_sam_predictor = SamPredictor(_sam_model)
|
||||
|
||||
print(f"SAM model loaded on {device}")
|
||||
return _sam_predictor
|
||||
|
||||
|
||||
def _sam_select_local(img_array: np.ndarray, x: int, y: int) -> np.ndarray:
|
||||
"""Use local SAM model for selection (no API calls, runs offline)"""
|
||||
predictor = _get_sam_model()
|
||||
|
||||
# Set image
|
||||
predictor.set_image(img_array)
|
||||
|
||||
# Point coordinates (x, y) and label (1 = foreground)
|
||||
input_point = np.array([[x, y]])
|
||||
input_label = np.array([1])
|
||||
|
||||
# Get mask prediction
|
||||
masks, scores, _ = predictor.predict(
|
||||
point_coords=input_point,
|
||||
point_labels=input_label,
|
||||
multimask_output=True, # Get multiple mask options
|
||||
)
|
||||
|
||||
# Use the mask with highest score
|
||||
best_mask_idx = np.argmax(scores)
|
||||
mask = masks[best_mask_idx]
|
||||
|
||||
return mask.astype(np.uint8)
|
||||
|
||||
|
||||
async def _sam_select(img_array: np.ndarray, x: int, y: int) -> np.ndarray:
|
||||
"""Use Segment Anything Model for selection via Replicate API"""
|
||||
"""
|
||||
Smart object selection using SAM (Segment Anything Model).
|
||||
|
||||
Priority:
|
||||
1. Local SAM model (free, fast, offline)
|
||||
2. Replicate API (if local not available and API key set)
|
||||
3. Raises exception if neither available
|
||||
"""
|
||||
# Try local SAM first (free, no API calls)
|
||||
try:
|
||||
return _sam_select_local(img_array, x, y)
|
||||
except FileNotFoundError as e:
|
||||
print(f"Local SAM not available: {e}")
|
||||
except ImportError as e:
|
||||
print(f"SAM dependencies not installed: {e}")
|
||||
except Exception as e:
|
||||
print(f"Local SAM failed: {e}")
|
||||
|
||||
# Fall back to Replicate API
|
||||
from app.config import settings
|
||||
|
||||
if not settings.replicate_api_key:
|
||||
raise ValueError(
|
||||
"SAM model not available. Either:\n"
|
||||
" 1. Download local model: docker exec -it ai-photo-edit-backend python /scripts/download_sam_model.py\n"
|
||||
" 2. Or set REPLICATE_API_KEY in .env for cloud SAM"
|
||||
)
|
||||
|
||||
return await _sam_select_replicate(img_array, x, y)
|
||||
|
||||
|
||||
async def _sam_select_replicate(img_array: np.ndarray, x: int, y: int) -> np.ndarray:
|
||||
"""Fallback: Use SAM via Replicate API (requires API key, costs ~$0.002/call)"""
|
||||
import httpx
|
||||
import base64
|
||||
import asyncio
|
||||
from app.config import settings
|
||||
|
||||
if not settings.replicate_api_key:
|
||||
raise ValueError("REPLICATE_API_KEY not configured")
|
||||
|
||||
# Convert image to base64
|
||||
img = Image.fromarray(img_array)
|
||||
buffer = BytesIO()
|
||||
img.save(buffer, format='PNG')
|
||||
img_b64 = base64.b64encode(buffer.getvalue()).decode('utf-8')
|
||||
|
||||
# Call SAM via Replicate
|
||||
async with httpx.AsyncClient(timeout=120.0) as client:
|
||||
# Use SAM model on Replicate
|
||||
prediction_data = {
|
||||
"version": "meta/sam-2-image:fe97b453d6525baeeb530595c74a3c4f567c1f655ee2a0fee11f76bd1d31e495",
|
||||
"input": {
|
||||
"image": f"data:image/png;base64,{img_b64}",
|
||||
"point_coords": f"{x},{y}",
|
||||
"point_labels": "1", # 1 = foreground point
|
||||
"point_labels": "1",
|
||||
}
|
||||
}
|
||||
|
||||
@@ -205,7 +312,6 @@ async def _sam_select(img_array: np.ndarray, x: int, y: int) -> np.ndarray:
|
||||
'Content-Type': 'application/json'
|
||||
}
|
||||
|
||||
# Start prediction
|
||||
response = await client.post(
|
||||
"https://api.replicate.com/v1/predictions",
|
||||
json=prediction_data,
|
||||
@@ -216,20 +322,15 @@ async def _sam_select(img_array: np.ndarray, x: int, y: int) -> np.ndarray:
|
||||
raise Exception(f"Replicate API error: {response.text}")
|
||||
|
||||
prediction = response.json()
|
||||
prediction_url = prediction['urls']['get']
|
||||
|
||||
# Poll for completion
|
||||
prediction_url = prediction['urls']['get']
|
||||
max_attempts = 60
|
||||
attempt = 0
|
||||
|
||||
while attempt < max_attempts:
|
||||
for _ in range(60):
|
||||
await asyncio.sleep(2)
|
||||
|
||||
status_response = await client.get(prediction_url, headers=headers)
|
||||
status_data = status_response.json()
|
||||
|
||||
if status_data['status'] == 'succeeded':
|
||||
# Download mask image
|
||||
mask_url = status_data['output']
|
||||
if isinstance(mask_url, list):
|
||||
mask_url = mask_url[0]
|
||||
@@ -237,20 +338,17 @@ async def _sam_select(img_array: np.ndarray, x: int, y: int) -> np.ndarray:
|
||||
mask_response = await client.get(mask_url)
|
||||
mask_img = Image.open(BytesIO(mask_response.content)).convert('L')
|
||||
|
||||
# Resize if needed
|
||||
if mask_img.size != (img_array.shape[1], img_array.shape[0]):
|
||||
mask_img = mask_img.resize(
|
||||
(img_array.shape[1], img_array.shape[0]),
|
||||
Image.Resampling.LANCZOS
|
||||
)
|
||||
|
||||
return np.array(mask_img) // 255 # Normalize to 0-1
|
||||
return np.array(mask_img) // 255
|
||||
|
||||
elif status_data['status'] == 'failed':
|
||||
raise Exception(f"SAM prediction failed: {status_data.get('error')}")
|
||||
|
||||
attempt += 1
|
||||
|
||||
raise Exception("SAM prediction timed out")
|
||||
|
||||
|
||||
|
||||
@@ -0,0 +1,61 @@
|
||||
#!/bin/bash
|
||||
# =============================================================================
|
||||
# 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
|
||||
# =============================================================================
|
||||
|
||||
set -e
|
||||
|
||||
echo "=========================================="
|
||||
echo "AI Photo Edit - Starting Up"
|
||||
echo "=========================================="
|
||||
|
||||
# Ensure data directories exist
|
||||
mkdir -p /app/data/projects
|
||||
mkdir -p /app/data/patches
|
||||
mkdir -p /app/data/models
|
||||
|
||||
# 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
|
||||
|
||||
echo ""
|
||||
echo "Checking SAM model (Smart Select)..."
|
||||
echo "------------------------------------------"
|
||||
if [ -f "/app/data/models/sam_model.pth" ] || \
|
||||
[ -f "/app/data/models/sam_vit_b_01ec64.pth" ] || \
|
||||
[ -f "/app/data/models/sam_vit_l_0b3195.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 ""
|
||||
fi
|
||||
|
||||
echo ""
|
||||
echo "=========================================="
|
||||
echo "Starting FastAPI server..."
|
||||
echo "=========================================="
|
||||
|
||||
# Start the server
|
||||
exec uvicorn app.main:app --host 0.0.0.0 --port 8000
|
||||
@@ -16,3 +16,7 @@ opencv-python-headless==4.9.0.80
|
||||
scikit-image==0.22.0
|
||||
rembg==2.0.50
|
||||
onnxruntime==1.16.3
|
||||
# SAM (Segment Anything) for smart object selection - runs locally, no API needed
|
||||
torch==2.1.2
|
||||
torchvision==0.16.2
|
||||
segment-anything @ git+https://github.com/facebookresearch/segment-anything.git
|
||||
|
||||
@@ -0,0 +1,126 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Download SAM (Segment Anything Model) for local inference.
|
||||
|
||||
This script downloads the SAM model checkpoint to a persistent directory
|
||||
so it survives container rebuilds.
|
||||
|
||||
Models available:
|
||||
- sam_vit_b: ~375MB (default, good balance of speed/quality)
|
||||
- sam_vit_l: ~1.2GB (better quality, slower)
|
||||
- sam_vit_h: ~2.5GB (best quality, slowest)
|
||||
|
||||
Usage:
|
||||
python scripts/download_sam_model.py [model_type]
|
||||
|
||||
model_type: vit_b (default), vit_l, or vit_h
|
||||
"""
|
||||
|
||||
import os
|
||||
import sys
|
||||
import urllib.request
|
||||
from pathlib import Path
|
||||
|
||||
# Model URLs from Meta's official releases
|
||||
SAM_MODELS = {
|
||||
'vit_b': {
|
||||
'url': 'https://dl.fbaipublicfiles.com/segment_anything/sam_vit_b_01ec64.pth',
|
||||
'filename': 'sam_vit_b_01ec64.pth',
|
||||
'size': '375MB'
|
||||
},
|
||||
'vit_l': {
|
||||
'url': 'https://dl.fbaipublicfiles.com/segment_anything/sam_vit_l_0b3195.pth',
|
||||
'filename': 'sam_vit_l_0b3195.pth',
|
||||
'size': '1.2GB'
|
||||
},
|
||||
'vit_h': {
|
||||
'url': 'https://dl.fbaipublicfiles.com/segment_anything/sam_vit_h_4b8939.pth',
|
||||
'filename': 'sam_vit_h_4b8939.pth',
|
||||
'size': '2.5GB'
|
||||
}
|
||||
}
|
||||
|
||||
def download_with_progress(url: str, dest_path: Path):
|
||||
"""Download file with progress indicator"""
|
||||
print(f"Downloading to: {dest_path}")
|
||||
|
||||
def progress_hook(count, block_size, total_size):
|
||||
percent = int(count * block_size * 100 / total_size)
|
||||
mb_done = count * block_size / (1024 * 1024)
|
||||
mb_total = total_size / (1024 * 1024)
|
||||
sys.stdout.write(f"\r Progress: {percent}% ({mb_done:.1f}/{mb_total:.1f} MB)")
|
||||
sys.stdout.flush()
|
||||
|
||||
urllib.request.urlretrieve(url, dest_path, progress_hook)
|
||||
print("\n Download complete!")
|
||||
|
||||
def main():
|
||||
# Determine model type
|
||||
model_type = sys.argv[1] if len(sys.argv) > 1 else 'vit_b'
|
||||
|
||||
if model_type not in SAM_MODELS:
|
||||
print(f"Unknown model type: {model_type}")
|
||||
print(f"Available: {', '.join(SAM_MODELS.keys())}")
|
||||
sys.exit(1)
|
||||
|
||||
model_info = SAM_MODELS[model_type]
|
||||
|
||||
# Determine models directory
|
||||
# Check if running in Docker (mounted volume) or locally
|
||||
models_dir = Path('/app/data/models')
|
||||
if not models_dir.exists():
|
||||
models_dir = Path(__file__).parent.parent / 'data' / 'models'
|
||||
|
||||
models_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
dest_path = models_dir / model_info['filename']
|
||||
|
||||
print("=" * 60)
|
||||
print("SAM Model Downloader")
|
||||
print("=" * 60)
|
||||
print(f"Model: SAM {model_type.upper()}")
|
||||
print(f"Size: {model_info['size']}")
|
||||
print(f"License: Apache 2.0 (commercial use OK)")
|
||||
print("=" * 60)
|
||||
|
||||
# Check if already downloaded
|
||||
if dest_path.exists():
|
||||
print(f"\nModel already exists at: {dest_path}")
|
||||
print("To re-download, delete the file first.")
|
||||
|
||||
# Create symlink for easy access
|
||||
symlink_path = models_dir / 'sam_model.pth'
|
||||
if symlink_path.exists() or symlink_path.is_symlink():
|
||||
symlink_path.unlink()
|
||||
symlink_path.symlink_to(dest_path.name)
|
||||
print(f"Symlink created: {symlink_path} -> {dest_path.name}")
|
||||
return
|
||||
|
||||
print(f"\nDownloading SAM {model_type.upper()} ({model_info['size']})...")
|
||||
print("This is a one-time download. The model will persist across rebuilds.")
|
||||
print()
|
||||
|
||||
try:
|
||||
download_with_progress(model_info['url'], dest_path)
|
||||
|
||||
# Create symlink for easy access
|
||||
symlink_path = models_dir / 'sam_model.pth'
|
||||
if symlink_path.exists() or symlink_path.is_symlink():
|
||||
symlink_path.unlink()
|
||||
symlink_path.symlink_to(dest_path.name)
|
||||
|
||||
print()
|
||||
print("=" * 60)
|
||||
print("SUCCESS!")
|
||||
print(f"Model saved to: {dest_path}")
|
||||
print(f"Symlink: {symlink_path}")
|
||||
print("=" * 60)
|
||||
|
||||
except Exception as e:
|
||||
print(f"\nError downloading model: {e}")
|
||||
if dest_path.exists():
|
||||
dest_path.unlink()
|
||||
sys.exit(1)
|
||||
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
Reference in New Issue
Block a user