root/PaintPlus
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add local SAM model support for offline Smart Select

Browse Source 
- Add torch, torchvision, segment-anything to requirements
- Create download_sam_model.py script to fetch SAM checkpoint
- Update tools.py to use local SAM with Replicate API fallback
- Add SAM model check to entrypoint.sh with helpful instructions
- Model persists in /app/data/models via Docker volume mount
This commit is contained in:
Claude
2026-01-25 18:15:07 +00:00
parent 28e842e190
commit 8b96c86a94
4 changed files with 270 additions and 19 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/app/routers/tools.py
+117 -19
View File
@@ -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")
backend/entrypoint.sh
+23
View File
@@ -30,6 +30,29 @@ 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 "=========================================="
backend/requirements.txt
+4
View File
@@ -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
scripts/download_sam_model.py
+126
View File
@@ -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()