Update app.py
Browse files
app.py
CHANGED
|
@@ -1,6 +1,6 @@
|
|
| 1 |
#!/usr/bin/env python3
|
| 2 |
-
# app.py - Background Remover
|
| 3 |
-
#
|
| 4 |
|
| 5 |
import io
|
| 6 |
import time
|
|
@@ -15,27 +15,21 @@ import gc
|
|
| 15 |
import shutil
|
| 16 |
import zipfile
|
| 17 |
import json
|
| 18 |
-
import hashlib
|
| 19 |
import resource
|
| 20 |
import uuid
|
| 21 |
import base64
|
| 22 |
import urllib.parse
|
| 23 |
-
import subprocess
|
| 24 |
from collections import defaultdict, Counter, deque
|
| 25 |
from datetime import datetime
|
| 26 |
from pathlib import Path
|
| 27 |
-
from typing import Dict, Optional, List,
|
| 28 |
from concurrent.futures import ThreadPoolExecutor
|
| 29 |
|
| 30 |
import psutil
|
| 31 |
|
| 32 |
# ============================================================
|
| 33 |
-
# ONNX RUNTIME
|
| 34 |
# ============================================================
|
| 35 |
-
# We previously had enable_cpu_mem_arena=False which caused every
|
| 36 |
-
# intermediate tensor to allocate fresh. silueta's int8-quantized
|
| 37 |
-
# Conv nodes need ~118MB per tensor and this caused OOM storms.
|
| 38 |
-
# Arena ON lets ORT reuse buffers between inferences.
|
| 39 |
try:
|
| 40 |
import onnxruntime as _ort
|
| 41 |
_orig_session_init = _ort.InferenceSession.__init__
|
|
@@ -46,8 +40,8 @@ try:
|
|
| 46 |
so = _ort.SessionOptions()
|
| 47 |
so.intra_op_num_threads = 2
|
| 48 |
so.inter_op_num_threads = 1
|
| 49 |
-
so.enable_cpu_mem_arena = True
|
| 50 |
-
so.enable_mem_pattern = True
|
| 51 |
so.graph_optimization_level = _ort.GraphOptimizationLevel.ORT_ENABLE_ALL
|
| 52 |
kwargs['sess_options'] = so
|
| 53 |
except Exception:
|
|
@@ -55,13 +49,13 @@ try:
|
|
| 55 |
return _orig_session_init(self, *args, **kwargs)
|
| 56 |
|
| 57 |
_ort.InferenceSession.__init__ = _patched_session_init
|
| 58 |
-
print("
|
| 59 |
except Exception as e:
|
| 60 |
-
print(f"
|
| 61 |
|
| 62 |
import numpy as np
|
| 63 |
-
from fastapi import FastAPI, File, UploadFile, HTTPException, Query,
|
| 64 |
-
from fastapi.responses import StreamingResponse, JSONResponse,
|
| 65 |
from fastapi.middleware.cors import CORSMiddleware
|
| 66 |
from PIL import Image, ImageDraw
|
| 67 |
import cv2
|
|
@@ -79,7 +73,7 @@ logger = logging.getLogger(__name__)
|
|
| 79 |
|
| 80 |
|
| 81 |
# ============================================================
|
| 82 |
-
#
|
| 83 |
# ============================================================
|
| 84 |
class ResolutionManager:
|
| 85 |
RECOMMENDED = {
|
|
@@ -141,7 +135,7 @@ def get_container_memory() -> Dict:
|
|
| 141 |
|
| 142 |
|
| 143 |
class ModelManager:
|
| 144 |
-
def __init__(self, max_models: int =
|
| 145 |
self.models: Dict[str, Any] = {}
|
| 146 |
self.load_times: Dict[str, float] = {}
|
| 147 |
self.use_counts: Dict[str, int] = {}
|
|
@@ -155,7 +149,7 @@ class ModelManager:
|
|
| 155 |
try:
|
| 156 |
mem = get_container_memory()
|
| 157 |
if 'cgroup' in mem.get('source', '') and mem['percent'] > 82:
|
| 158 |
-
logger.warning(f"
|
| 159 |
self._evict_models()
|
| 160 |
gc.collect()
|
| 161 |
except Exception:
|
|
@@ -171,7 +165,7 @@ class ModelManager:
|
|
| 171 |
gc.collect()
|
| 172 |
try:
|
| 173 |
mem = get_container_memory()
|
| 174 |
-
logger.info(f"
|
| 175 |
f"({mem['percent']:.1f}% of {mem['total_mb']}MB)")
|
| 176 |
except Exception:
|
| 177 |
pass
|
|
@@ -187,7 +181,7 @@ class ModelManager:
|
|
| 187 |
if len(self.models) <= 1:
|
| 188 |
return
|
| 189 |
to_remove = min(self.models.keys(), key=lambda x: self.use_counts.get(x, 0))
|
| 190 |
-
logger.info(f"
|
| 191 |
del self.models[to_remove]
|
| 192 |
self.use_counts.pop(to_remove, None)
|
| 193 |
gc.collect()
|
|
@@ -421,16 +415,13 @@ class ContentClassifier:
|
|
| 421 |
return 'complex'
|
| 422 |
|
| 423 |
|
| 424 |
-
# ============================================================
|
| 425 |
-
# MEMORY + ENV
|
| 426 |
-
# ============================================================
|
| 427 |
def limit_memory():
|
| 428 |
try:
|
| 429 |
soft, hard = resource.getrlimit(resource.RLIMIT_AS)
|
| 430 |
resource.setrlimit(resource.RLIMIT_AS, (12 * 1024 * 1024 * 1024, hard))
|
| 431 |
-
logger.info("
|
| 432 |
except Exception as e:
|
| 433 |
-
logger.warning(f"
|
| 434 |
|
| 435 |
|
| 436 |
limit_memory()
|
|
@@ -448,12 +439,9 @@ for d in ["/tmp/u2net_models", "/tmp/uploads", "/tmp/outputs", "/tmp/frames"]:
|
|
| 448 |
cpu_executor = ThreadPoolExecutor(max_workers=2)
|
| 449 |
|
| 450 |
|
| 451 |
-
# ============================================================
|
| 452 |
-
# SINGLETONS
|
| 453 |
-
# ============================================================
|
| 454 |
content_classifier = ContentClassifier()
|
| 455 |
resolution_manager = ResolutionManager()
|
| 456 |
-
model_manager = ModelManager(max_models=
|
| 457 |
mask_checker = MaskQualityChecker()
|
| 458 |
hair_refiner = HairRefiner()
|
| 459 |
graphic_processor = GraphicProcessor()
|
|
@@ -462,13 +450,14 @@ graphic_processor = GraphicProcessor()
|
|
| 462 |
# ============================================================
|
| 463 |
# MODEL CONFIG
|
| 464 |
# ============================================================
|
|
|
|
| 465 |
PRELOAD_MODELS = [
|
| 466 |
("u2netp", "~4.7 MB disk / ~180 MB RAM"),
|
| 467 |
-
("isnet-general-use", "~176 MB disk / ~800 MB RAM"),
|
| 468 |
("birefnet-general-lite", "~12 MB disk / ~250 MB RAM"),
|
| 469 |
]
|
| 470 |
|
| 471 |
LAZY_MODELS = {
|
|
|
|
| 472 |
"silueta": "~43 MB disk / ~200 MB RAM",
|
| 473 |
"u2net": "~176 MB disk / ~800 MB RAM",
|
| 474 |
"u2net_human_seg": "~176 MB disk / ~800 MB RAM",
|
|
@@ -528,11 +517,6 @@ MODEL_RES_CAPS = {
|
|
| 528 |
|
| 529 |
CLASSICAL_PROCESSORS = {'flood_fill', 'graphic_classical'}
|
| 530 |
DEFAULT_MODEL_NAME = 'isnet-general-use'
|
| 531 |
-
|
| 532 |
-
# ⬅ ADDED: default model for the video frame-batch path.
|
| 533 |
-
# silueta was causing OOM storms on free-tier HF and silently
|
| 534 |
-
# returning transparent frames. birefnet-general-lite is stable
|
| 535 |
-
# and pre-loaded.
|
| 536 |
DEFAULT_VIDEO_MODEL = 'birefnet-general-lite'
|
| 537 |
|
| 538 |
DISPLAY_NAMES = {
|
|
@@ -543,7 +527,7 @@ DISPLAY_NAMES = {
|
|
| 543 |
'u2net_human_seg': 'U2-Net Human',
|
| 544 |
'u2net_cloth_seg': 'U2-Net Cloth',
|
| 545 |
'silueta': 'Silueta',
|
| 546 |
-
'isnet-general-use': '
|
| 547 |
'isnet-anime': 'ISNet Anime',
|
| 548 |
'birefnet-general': 'BiRefNet',
|
| 549 |
'birefnet-general-lite': 'BiRefNet Lite',
|
|
@@ -551,18 +535,15 @@ DISPLAY_NAMES = {
|
|
| 551 |
}
|
| 552 |
|
| 553 |
|
| 554 |
-
# ============================================================
|
| 555 |
-
# SESSION LOADER
|
| 556 |
-
# ============================================================
|
| 557 |
def _load_session(name: str):
|
| 558 |
from rembg import new_session
|
| 559 |
resolved = MODEL_ALIASES.get(name, name)
|
| 560 |
if resolved not in VALID_REMBG_MODELS:
|
| 561 |
-
logger.warning(f"
|
| 562 |
resolved = 'isnet-general-use'
|
| 563 |
if resolved != name:
|
| 564 |
-
logger.info(f"
|
| 565 |
-
logger.info(f"
|
| 566 |
return new_session(resolved)
|
| 567 |
|
| 568 |
|
|
@@ -683,9 +664,6 @@ def process_with_neural(image, session, model_name, force_no_matting=False):
|
|
| 683 |
return None, diagnostics
|
| 684 |
|
| 685 |
|
| 686 |
-
# ============================================================
|
| 687 |
-
# CLASSICAL FALLBACK
|
| 688 |
-
# ============================================================
|
| 689 |
def process_graphic_classical(image):
|
| 690 |
try:
|
| 691 |
if image.mode != 'RGBA':
|
|
@@ -757,24 +735,24 @@ default_model_name = DEFAULT_MODEL_NAME
|
|
| 757 |
PYMATTING_AVAILABLE = False
|
| 758 |
try:
|
| 759 |
from pymatting import cutout, estimate_foreground_ml, estimate_alpha_cf
|
| 760 |
-
from pymatting.util.util import stack_images
|
| 761 |
PYMATTING_AVAILABLE = True
|
| 762 |
-
logger.info("
|
| 763 |
except ImportError as e:
|
| 764 |
-
logger.warning(f"
|
| 765 |
|
| 766 |
|
| 767 |
def pre_download_models():
|
| 768 |
-
logger.info(f"
|
| 769 |
for name, size in PRELOAD_MODELS:
|
| 770 |
try:
|
| 771 |
logger.info(f" Downloading {name} ({size})...")
|
| 772 |
session = model_manager.get_model(name, _load_session)
|
| 773 |
if session is not None:
|
| 774 |
-
logger.info(f"
|
| 775 |
except Exception as e:
|
| 776 |
-
logger.warning(f"
|
| 777 |
-
logger.info(
|
| 778 |
|
| 779 |
|
| 780 |
def initialize_models():
|
|
@@ -795,7 +773,7 @@ def initialize_models():
|
|
| 795 |
default_model_name = loaded[0] if loaded else 'u2netp'
|
| 796 |
AI_AVAILABLE = len(loaded) > 0
|
| 797 |
if AI_AVAILABLE:
|
| 798 |
-
logger.info(f"
|
| 799 |
except Exception as e:
|
| 800 |
logger.error(f"AI setup failed: {e}")
|
| 801 |
AI_AVAILABLE = False
|
|
@@ -805,7 +783,7 @@ def _prewarm_models():
|
|
| 805 |
global remove_func
|
| 806 |
if remove_func is None:
|
| 807 |
return
|
| 808 |
-
logger.info("
|
| 809 |
try:
|
| 810 |
test_img = Image.new('RGB', (64, 64), (255, 255, 255))
|
| 811 |
ImageDraw.Draw(test_img).rectangle([20, 20, 44, 44], fill=(220, 40, 40))
|
|
@@ -815,36 +793,33 @@ def _prewarm_models():
|
|
| 815 |
continue
|
| 816 |
try:
|
| 817 |
t0 = time.time()
|
| 818 |
-
remove_func(test_img, session=session, alpha_matting=False)
|
| 819 |
-
t1 = time.time()
|
| 820 |
result = remove_func(test_img, session=session, alpha_matting=False)
|
| 821 |
-
|
| 822 |
if result.mode != 'RGBA':
|
| 823 |
result = result.convert('RGBA')
|
| 824 |
arr = np.array(result)
|
| 825 |
cy, cx = arr.shape[0] // 2, arr.shape[1] // 2
|
| 826 |
ca = int(arr[cy, cx, 3]); corner = int(arr[2, 2, 3])
|
| 827 |
if ca > 200 and corner < 50:
|
| 828 |
-
logger.info(f"
|
| 829 |
-
f"first={t1-t0:.2f}s warm={t2-t1:.2f}s")
|
| 830 |
else:
|
| 831 |
-
logger.warning(f"
|
| 832 |
except Exception as e:
|
| 833 |
-
logger.warning(f"
|
| 834 |
except Exception as e:
|
| 835 |
logger.warning(f"Pre-warm skipped: {e}")
|
| 836 |
|
| 837 |
|
| 838 |
logger.info("=" * 60)
|
| 839 |
-
logger.info("
|
| 840 |
logger.info("=" * 60)
|
| 841 |
pre_download_models()
|
| 842 |
initialize_models()
|
| 843 |
_prewarm_models()
|
| 844 |
mem = get_container_memory()
|
| 845 |
-
logger.info(f"
|
| 846 |
-
logger.info(f"
|
| 847 |
-
logger.info(f"
|
| 848 |
logger.info("=" * 60)
|
| 849 |
|
| 850 |
|
|
@@ -869,9 +844,14 @@ def get_session_for_model(model_name):
|
|
| 869 |
|
| 870 |
|
| 871 |
# ============================================================
|
| 872 |
-
# ROUTER
|
| 873 |
# ============================================================
|
| 874 |
def detect_hair_complexity(image):
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 875 |
try:
|
| 876 |
gray = image.convert('L').resize((100, 100), Image.Resampling.LANCZOS)
|
| 877 |
np_img = np.array(gray, dtype=np.float32)
|
|
@@ -880,10 +860,12 @@ def detect_hair_complexity(image):
|
|
| 880 |
fine = np.sum(em > 10) / 10000
|
| 881 |
coarse = np.sum(em > 40) / 10000
|
| 882 |
fr = fine / (coarse + 0.001)
|
| 883 |
-
|
| 884 |
-
|
|
|
|
|
|
|
| 885 |
return {'complexity_score': score, 'route': 'wispy_hair'}
|
| 886 |
-
if score >
|
| 887 |
return {'complexity_score': score, 'route': 'complex_hair'}
|
| 888 |
return {'complexity_score': score, 'route': 'simple_hair'}
|
| 889 |
except Exception:
|
|
@@ -1000,7 +982,7 @@ def detect_fragments(image):
|
|
| 1000 |
# IMAGE PIPELINE
|
| 1001 |
# ============================================================
|
| 1002 |
def process_image_sync(image_data, model_name=None, transparent=True,
|
| 1003 |
-
max_size=
|
| 1004 |
use_alpha_matting=True, border_pixels=20, sharpen=True):
|
| 1005 |
req = str(uuid.uuid4())[:8]
|
| 1006 |
logger.info(f"[REQ {req}] === START ===")
|
|
@@ -1026,7 +1008,7 @@ def process_image_sync(image_data, model_name=None, transparent=True,
|
|
| 1026 |
'media_type': mt, 'extension': ext,
|
| 1027 |
'original_size': original_size, 'processed_size': image.size,
|
| 1028 |
'fragment_ratio': 0.0, 'quality_score': 100.0,
|
| 1029 |
-
'quality_label': 'Skipped', 'quality_emoji': '
|
| 1030 |
'has_fragments': False, 'issues': [],
|
| 1031 |
'needs_reprocess': False, 'router_enabled': False,
|
| 1032 |
'router_score': 0.0, 'router_reason': 'already transparent',
|
|
@@ -1102,10 +1084,10 @@ def process_image_sync(image_data, model_name=None, transparent=True,
|
|
| 1102 |
frag = detect_fragments(result)
|
| 1103 |
|
| 1104 |
qs = frag['quality_score']
|
| 1105 |
-
if qs >= 90: ql
|
| 1106 |
-
elif qs >= 75: ql
|
| 1107 |
-
elif qs >= 50: ql
|
| 1108 |
-
else: ql
|
| 1109 |
|
| 1110 |
buf = io.BytesIO()
|
| 1111 |
if transparent:
|
|
@@ -1124,7 +1106,7 @@ def process_image_sync(image_data, model_name=None, transparent=True,
|
|
| 1124 |
'original_size': original_size, 'processed_size': result.size,
|
| 1125 |
'fragment_ratio': frag['fragment_ratio'],
|
| 1126 |
'quality_score': frag['quality_score'],
|
| 1127 |
-
'quality_label': ql, 'quality_emoji':
|
| 1128 |
'has_fragments': frag['has_fragments'], 'issues': frag['issues'],
|
| 1129 |
'needs_reprocess': frag['has_fragments'],
|
| 1130 |
'router_enabled': router_enabled,
|
|
@@ -1143,22 +1125,10 @@ def process_image_sync(image_data, model_name=None, transparent=True,
|
|
| 1143 |
raise
|
| 1144 |
|
| 1145 |
|
| 1146 |
-
def cleanup_temp_files(*file_paths):
|
| 1147 |
-
for path in file_paths:
|
| 1148 |
-
if path and os.path.exists(path):
|
| 1149 |
-
try:
|
| 1150 |
-
if os.path.isdir(path):
|
| 1151 |
-
shutil.rmtree(path)
|
| 1152 |
-
else:
|
| 1153 |
-
os.unlink(path)
|
| 1154 |
-
except Exception as e:
|
| 1155 |
-
logger.warning(f"Cleanup failed for {path}: {e}")
|
| 1156 |
-
|
| 1157 |
-
|
| 1158 |
# ============================================================
|
| 1159 |
# FASTAPI
|
| 1160 |
# ============================================================
|
| 1161 |
-
app = FastAPI(title="Background Remover
|
| 1162 |
app.add_middleware(CORSMiddleware, allow_origins=["*"], allow_credentials=True,
|
| 1163 |
allow_methods=["*"], allow_headers=["*"])
|
| 1164 |
|
|
@@ -1177,7 +1147,7 @@ def _get_client_ip(request):
|
|
| 1177 |
@app.middleware("http")
|
| 1178 |
async def rate_limit_middleware(request, call_next):
|
| 1179 |
if request.url.path in ("/api/health", "/", "/docs", "/redoc", "/openapi.json",
|
| 1180 |
-
"/api/debug/env", "/api/debug/model-report"):
|
| 1181 |
return await call_next(request)
|
| 1182 |
cid = _get_client_ip(request)
|
| 1183 |
now = datetime.now()
|
|
@@ -1200,16 +1170,13 @@ def timeout_handler(seconds=90):
|
|
| 1200 |
return deco
|
| 1201 |
|
| 1202 |
|
| 1203 |
-
# ============================================================
|
| 1204 |
-
# IMAGE ENDPOINTS
|
| 1205 |
-
# ============================================================
|
| 1206 |
@app.post("/api/process/image")
|
| 1207 |
@timeout_handler(seconds=90)
|
| 1208 |
async def process_image_api(
|
| 1209 |
file: UploadFile = File(...),
|
| 1210 |
model: str = Query(None),
|
| 1211 |
transparent: bool = Query(True),
|
| 1212 |
-
max_size: int = Query(
|
| 1213 |
quality: int = Query(95, ge=70, le=100),
|
| 1214 |
preserve_hair: bool = Query(True),
|
| 1215 |
use_alpha_matting: bool = Query(True),
|
|
@@ -1242,7 +1209,6 @@ async def process_image_api(
|
|
| 1242 |
|
| 1243 |
pt = time.time() - t0
|
| 1244 |
fn = Path(file.filename).stem
|
| 1245 |
-
emoji = urllib.parse.quote(result_data.get('quality_emoji', ''), safe='')
|
| 1246 |
headers = {
|
| 1247 |
'X-Processing-Time': f'{pt:.3f}',
|
| 1248 |
'X-Model-Used': result_data['model'],
|
|
@@ -1250,7 +1216,6 @@ async def process_image_api(
|
|
| 1250 |
'X-Transparent': str(result_data['transparent']).lower(),
|
| 1251 |
'X-Quality-Score': f"{result_data['quality_score']:.1f}",
|
| 1252 |
'X-Quality-Label': result_data.get('quality_label', ''),
|
| 1253 |
-
'X-Quality-Emoji': emoji,
|
| 1254 |
'X-Fragment-Ratio': f"{result_data.get('fragment_ratio', 0):.3f}",
|
| 1255 |
'X-Needs-Reprocess': str(result_data.get('needs_reprocess', False)).lower(),
|
| 1256 |
'X-Router-Enabled': str(result_data.get('router_enabled', False)).lower(),
|
|
@@ -1269,59 +1234,6 @@ async def process_image_api(
|
|
| 1269 |
raise HTTPException(500, f"Processing failed: {str(e)[:200]}")
|
| 1270 |
|
| 1271 |
|
| 1272 |
-
@app.post("/api/process/image-png")
|
| 1273 |
-
@timeout_handler(seconds=90)
|
| 1274 |
-
async def process_image_png(
|
| 1275 |
-
file: UploadFile = File(...),
|
| 1276 |
-
model: str = Query(None),
|
| 1277 |
-
preserve_hair: bool = Query(True),
|
| 1278 |
-
use_alpha_matting: bool = Query(True),
|
| 1279 |
-
border_pixels: int = Query(20, ge=5, le=50),
|
| 1280 |
-
sharpen: bool = Query(True),
|
| 1281 |
-
):
|
| 1282 |
-
t0 = time.time()
|
| 1283 |
-
try:
|
| 1284 |
-
if not file.content_type or not file.content_type.startswith('image/'):
|
| 1285 |
-
raise HTTPException(400, "Not an image")
|
| 1286 |
-
contents = await file.read()
|
| 1287 |
-
if not contents: raise HTTPException(400, "Empty")
|
| 1288 |
-
if len(contents) > 30 * 1024 * 1024: raise HTTPException(400, "Too large")
|
| 1289 |
-
if not AI_AVAILABLE: raise HTTPException(503, "AI unavailable")
|
| 1290 |
-
user_model = None
|
| 1291 |
-
if model:
|
| 1292 |
-
resolved = MODEL_ALIASES.get(model, model)
|
| 1293 |
-
if resolved in VALID_REMBG_MODELS or resolved in CLASSICAL_PROCESSORS:
|
| 1294 |
-
user_model = resolved
|
| 1295 |
-
loop = asyncio.get_event_loop()
|
| 1296 |
-
pf = functools.partial(process_image_sync, image_data=contents,
|
| 1297 |
-
model_name=user_model, transparent=True,
|
| 1298 |
-
max_size=2048, quality=95,
|
| 1299 |
-
preserve_hair=preserve_hair,
|
| 1300 |
-
use_alpha_matting=use_alpha_matting,
|
| 1301 |
-
border_pixels=border_pixels, sharpen=sharpen)
|
| 1302 |
-
result_data = await loop.run_in_executor(cpu_executor, pf)
|
| 1303 |
-
pt = time.time() - t0
|
| 1304 |
-
fn = Path(file.filename).stem
|
| 1305 |
-
emoji = urllib.parse.quote(result_data.get('quality_emoji', ''), safe='')
|
| 1306 |
-
headers = {
|
| 1307 |
-
'X-Processing-Time': f'{pt:.3f}',
|
| 1308 |
-
'X-Model-Used': result_data['model'],
|
| 1309 |
-
'X-Model-Id': result_data['model_id'],
|
| 1310 |
-
'X-Transparent': 'true',
|
| 1311 |
-
'X-Quality-Score': f"{result_data['quality_score']:.1f}",
|
| 1312 |
-
'X-Quality-Label': result_data.get('quality_label', ''),
|
| 1313 |
-
'X-Quality-Emoji': emoji,
|
| 1314 |
-
'Content-Disposition': f"attachment; filename=processed_{fn}.png",
|
| 1315 |
-
}
|
| 1316 |
-
return StreamingResponse(io.BytesIO(result_data['data']),
|
| 1317 |
-
media_type='image/png', headers=headers)
|
| 1318 |
-
except HTTPException:
|
| 1319 |
-
raise
|
| 1320 |
-
except Exception as e:
|
| 1321 |
-
logger.error(f"PNG error: {e}")
|
| 1322 |
-
raise HTTPException(500, str(e)[:200])
|
| 1323 |
-
|
| 1324 |
-
|
| 1325 |
@app.post("/api/process/bulk")
|
| 1326 |
@timeout_handler(seconds=300)
|
| 1327 |
async def process_bulk_images(
|
|
@@ -1391,38 +1303,19 @@ async def process_bulk_images(
|
|
| 1391 |
raise HTTPException(500, str(e)[:200])
|
| 1392 |
|
| 1393 |
|
| 1394 |
-
@app.post("/api/process/bulk-png")
|
| 1395 |
-
@timeout_handler(seconds=300)
|
| 1396 |
-
async def process_bulk_png(
|
| 1397 |
-
files: List[UploadFile] = File(...),
|
| 1398 |
-
model: str = Query(None),
|
| 1399 |
-
preserve_hair: bool = Query(True),
|
| 1400 |
-
use_alpha_matting: bool = Query(True),
|
| 1401 |
-
border_pixels: int = Query(20, ge=5, le=50),
|
| 1402 |
-
sharpen: bool = Query(True),
|
| 1403 |
-
):
|
| 1404 |
-
return await process_bulk_images(files=files, model=model, transparent=True,
|
| 1405 |
-
max_size=2048, preserve_hair=preserve_hair,
|
| 1406 |
-
use_alpha_matting=use_alpha_matting,
|
| 1407 |
-
border_pixels=border_pixels, sharpen=sharpen)
|
| 1408 |
-
|
| 1409 |
-
|
| 1410 |
# ============================================================
|
| 1411 |
-
# FRAME-BATCH ENDPOINT
|
| 1412 |
# ============================================================
|
| 1413 |
-
# Key fixes over v31.4.0:
|
| 1414 |
-
# 1. Default model changed to birefnet-general-lite (silueta was
|
| 1415 |
-
# OOMing on free-tier HF and producing blank frames).
|
| 1416 |
-
# 2. Detects fully-transparent AND fully-opaque frames (silent
|
| 1417 |
-
# failures) and falls through a model chain instead of shipping
|
| 1418 |
-
# garbage to the client.
|
| 1419 |
-
# 3. Returns the actual model used per batch.
|
| 1420 |
@app.post("/api/process/frames-batch")
|
| 1421 |
async def process_frames_batch(
|
| 1422 |
files: List[UploadFile] = File(...),
|
| 1423 |
model: str = Query(DEFAULT_VIDEO_MODEL),
|
| 1424 |
transparent: bool = Query(True),
|
| 1425 |
):
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1426 |
resolved = MODEL_ALIASES.get(model, model)
|
| 1427 |
if resolved not in VALID_REMBG_MODELS:
|
| 1428 |
resolved = DEFAULT_VIDEO_MODEL
|
|
@@ -1431,8 +1324,6 @@ async def process_frames_batch(
|
|
| 1431 |
if session is None:
|
| 1432 |
raise HTTPException(503, f"No session for {resolved}")
|
| 1433 |
|
| 1434 |
-
# Chain of models to try in order. First is user-selected, then
|
| 1435 |
-
# stable fallbacks.
|
| 1436 |
fallback_chain = [resolved]
|
| 1437 |
for candidate in ('birefnet-general-lite', 'isnet-general-use', 'u2netp'):
|
| 1438 |
if candidate not in fallback_chain:
|
|
@@ -1468,20 +1359,13 @@ async def process_frames_batch(
|
|
| 1468 |
min_a = int(alpha.min())
|
| 1469 |
max_a = int(alpha.max())
|
| 1470 |
|
| 1471 |
-
# Reject fully-transparent (silent fail)
|
| 1472 |
if max_a < 10:
|
| 1473 |
-
logger.warning(
|
| 1474 |
-
|
| 1475 |
-
f"frame {idx} trying next"
|
| 1476 |
-
)
|
| 1477 |
continue
|
| 1478 |
-
|
| 1479 |
-
# Reject fully-opaque (no-op / model didn't segment)
|
| 1480 |
if min_a > 245:
|
| 1481 |
-
logger.warning(
|
| 1482 |
-
|
| 1483 |
-
f"frame {idx} trying next"
|
| 1484 |
-
)
|
| 1485 |
continue
|
| 1486 |
|
| 1487 |
result = res
|
|
@@ -1489,8 +1373,6 @@ async def process_frames_batch(
|
|
| 1489 |
break
|
| 1490 |
|
| 1491 |
if result is None:
|
| 1492 |
-
# Every model failed on this frame. Send the original
|
| 1493 |
-
# through, flagged so the client can count fallbacks.
|
| 1494 |
logger.error(f"[frames-batch] frame {idx}: all models failed")
|
| 1495 |
buf = io.BytesIO()
|
| 1496 |
img.convert("RGBA").save(buf, format="PNG", compress_level=3)
|
|
@@ -1517,18 +1399,10 @@ async def process_frames_batch(
|
|
| 1517 |
})
|
| 1518 |
except Exception as e:
|
| 1519 |
logger.error(f"[frames-batch] frame {idx} failed: {e}")
|
| 1520 |
-
results.append({
|
| 1521 |
-
"data": "",
|
| 1522 |
-
"format": "png",
|
| 1523 |
-
"error": str(e)[:120],
|
| 1524 |
-
})
|
| 1525 |
|
| 1526 |
used = next((r.get("model") for r in results if r.get("model")), resolved)
|
| 1527 |
-
return {
|
| 1528 |
-
"frames": results,
|
| 1529 |
-
"count": len(results),
|
| 1530 |
-
"model": used,
|
| 1531 |
-
}
|
| 1532 |
|
| 1533 |
|
| 1534 |
# ============================================================
|
|
@@ -1607,57 +1481,6 @@ async def editor_inpaint(
|
|
| 1607 |
raise HTTPException(500, f"Inpainting failed: {str(e)[:200]}")
|
| 1608 |
|
| 1609 |
|
| 1610 |
-
@app.post("/api/editor/compose")
|
| 1611 |
-
async def editor_compose(
|
| 1612 |
-
foreground: UploadFile = File(...),
|
| 1613 |
-
background: UploadFile = File(None),
|
| 1614 |
-
bg_color: str = Form('#FFFFFF'),
|
| 1615 |
-
bg_gradient_end: str = Form(None),
|
| 1616 |
-
fg_offset_x: int = Form(0),
|
| 1617 |
-
fg_offset_y: int = Form(0),
|
| 1618 |
-
fg_scale: float = Form(1.0),
|
| 1619 |
-
canvas_w: int = Form(0),
|
| 1620 |
-
canvas_h: int = Form(0),
|
| 1621 |
-
):
|
| 1622 |
-
try:
|
| 1623 |
-
fg_contents = await foreground.read()
|
| 1624 |
-
fg = Image.open(io.BytesIO(fg_contents)).convert('RGBA')
|
| 1625 |
-
if canvas_w > 0 and canvas_h > 0:
|
| 1626 |
-
cw, ch = canvas_w, canvas_h
|
| 1627 |
-
else:
|
| 1628 |
-
cw, ch = fg.size
|
| 1629 |
-
if fg_scale != 1.0:
|
| 1630 |
-
fg = fg.resize((max(1, int(fg.width * fg_scale)),
|
| 1631 |
-
max(1, int(fg.height * fg_scale))), Image.Resampling.LANCZOS)
|
| 1632 |
-
if background is not None:
|
| 1633 |
-
bg_contents = await background.read()
|
| 1634 |
-
bg_img = Image.open(io.BytesIO(bg_contents)).convert('RGBA')
|
| 1635 |
-
scale = max(cw / bg_img.width, ch / bg_img.height)
|
| 1636 |
-
bg_img = bg_img.resize((int(bg_img.width * scale), int(bg_img.height * scale)),
|
| 1637 |
-
Image.Resampling.LANCZOS)
|
| 1638 |
-
left = (bg_img.width - cw) // 2
|
| 1639 |
-
top = (bg_img.height - ch) // 2
|
| 1640 |
-
canvas = bg_img.crop((left, top, left + cw, top + ch))
|
| 1641 |
-
elif bg_gradient_end:
|
| 1642 |
-
c1 = _hex_to_rgb(bg_color); c2 = _hex_to_rgb(bg_gradient_end)
|
| 1643 |
-
arr = np.zeros((ch, cw, 3), dtype=np.uint8)
|
| 1644 |
-
for y in range(ch):
|
| 1645 |
-
t = y / max(1, ch - 1)
|
| 1646 |
-
arr[y, :] = [int(c1[i] * (1 - t) + c2[i] * t) for i in range(3)]
|
| 1647 |
-
canvas = Image.fromarray(arr, 'RGB').convert('RGBA')
|
| 1648 |
-
else:
|
| 1649 |
-
c = _hex_to_rgb(bg_color)
|
| 1650 |
-
canvas = Image.new('RGBA', (cw, ch), c + (255,))
|
| 1651 |
-
canvas.paste(fg, (fg_offset_x, fg_offset_y), fg)
|
| 1652 |
-
buf = io.BytesIO()
|
| 1653 |
-
canvas.save(buf, format='PNG', optimize=True)
|
| 1654 |
-
buf.seek(0)
|
| 1655 |
-
return StreamingResponse(buf, media_type='image/png')
|
| 1656 |
-
except Exception as e:
|
| 1657 |
-
logger.error(f"compose failed: {e}")
|
| 1658 |
-
raise HTTPException(500, f"Compose failed: {str(e)[:200]}")
|
| 1659 |
-
|
| 1660 |
-
|
| 1661 |
def _hex_to_rgb(hex_str):
|
| 1662 |
h = hex_str.lstrip('#')
|
| 1663 |
if len(h) == 3:
|
|
@@ -1666,8 +1489,14 @@ def _hex_to_rgb(hex_str):
|
|
| 1666 |
|
| 1667 |
|
| 1668 |
# ============================================================
|
| 1669 |
-
# DEBUG / HEALTH
|
| 1670 |
# ============================================================
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1671 |
@app.get("/api/debug/env")
|
| 1672 |
async def debug_env():
|
| 1673 |
m = get_container_memory()
|
|
@@ -1686,69 +1515,25 @@ async def debug_env():
|
|
| 1686 |
'matting_models': sorted(MATTING_MODELS)}
|
| 1687 |
|
| 1688 |
|
| 1689 |
-
@app.get("/api/debug/model-report")
|
| 1690 |
-
async def model_report():
|
| 1691 |
-
if remove_func is None:
|
| 1692 |
-
return JSONResponse({'error': 'remove_func not ready'}, status_code=503)
|
| 1693 |
-
test_img = Image.new('RGB', (128, 128), (255, 255, 255))
|
| 1694 |
-
ImageDraw.Draw(test_img).rectangle([40, 40, 88, 88], fill=(220, 40, 40))
|
| 1695 |
-
report = {}
|
| 1696 |
-
for name in sorted(VALID_REMBG_MODELS):
|
| 1697 |
-
entry = {'model': name}
|
| 1698 |
-
try:
|
| 1699 |
-
session = get_session_for_model(name)
|
| 1700 |
-
if session is None:
|
| 1701 |
-
entry['error'] = 'session-load-failed'
|
| 1702 |
-
report[name] = entry
|
| 1703 |
-
continue
|
| 1704 |
-
t0 = time.time()
|
| 1705 |
-
try:
|
| 1706 |
-
result = remove_func(test_img, session=session, alpha_matting=False)
|
| 1707 |
-
dt = time.time() - t0
|
| 1708 |
-
if result.mode != 'RGBA':
|
| 1709 |
-
result = result.convert('RGBA')
|
| 1710 |
-
arr = np.array(result)
|
| 1711 |
-
stats = {'min': int(arr[:, :, 3].min()), 'max': int(arr[:, :, 3].max()),
|
| 1712 |
-
'mean': round(float(arr[:, :, 3].mean()), 2),
|
| 1713 |
-
'center_alpha': int(arr[64, 64, 3]),
|
| 1714 |
-
'corner_alpha': int(arr[5, 5, 3])}
|
| 1715 |
-
noop, reason = _is_noop(test_img, result)
|
| 1716 |
-
verdict = 'OK' if (stats['center_alpha'] > 200 and stats['corner_alpha'] < 50) else 'BROKEN'
|
| 1717 |
-
entry.update({'elapsed_s': round(dt, 3), 'stats': stats,
|
| 1718 |
-
'is_noop': noop, 'reason': reason, 'verdict': verdict})
|
| 1719 |
-
except Exception as e:
|
| 1720 |
-
entry['error'] = f'{type(e).__name__}: {str(e)[:200]}'
|
| 1721 |
-
except Exception as e:
|
| 1722 |
-
entry['error'] = f'{type(e).__name__}: {str(e)[:200]}'
|
| 1723 |
-
report[name] = entry
|
| 1724 |
-
return JSONResponse({
|
| 1725 |
-
'generated_at': datetime.now().isoformat(),
|
| 1726 |
-
'models': report,
|
| 1727 |
-
'summary': {'total': len(report),
|
| 1728 |
-
'ok': sum(1 for v in report.values() if v.get('verdict') == 'OK'),
|
| 1729 |
-
'broken': sum(1 for v in report.values() if v.get('verdict') == 'BROKEN'),
|
| 1730 |
-
'errored': sum(1 for v in report.values() if 'error' in v)}})
|
| 1731 |
-
|
| 1732 |
-
|
| 1733 |
@app.get("/")
|
| 1734 |
async def home():
|
| 1735 |
stats = model_manager.get_stats()
|
| 1736 |
mem = get_container_memory()
|
| 1737 |
return HTMLResponse(f"""
|
| 1738 |
<!DOCTYPE html><html><head><meta charset="UTF-8">
|
| 1739 |
-
<title>Background Remover
|
| 1740 |
-
<style>body{{font-family:system-ui;background:
|
| 1741 |
color:#fff;min-height:100vh;display:flex;align-items:center;justify-content:center;
|
| 1742 |
-
margin:0;padding:20px}}.c{{max-width:900px;background:
|
| 1743 |
-
|
| 1744 |
-
|
| 1745 |
-
margin:20px 0;background:
|
| 1746 |
</head><body><div class="c">
|
| 1747 |
-
<h1>Background Remover
|
| 1748 |
-
<div class="s">AI Online
|
| 1749 |
RAM {mem['percent']:.0f}% ({mem['used_mb']}/{mem['total_mb']}MB)</div>
|
| 1750 |
<p>Video default: <strong>{DEFAULT_VIDEO_MODEL}</strong></p>
|
| 1751 |
-
<p><a href="/docs">API Docs</a>
|
| 1752 |
</div></body></html>
|
| 1753 |
""")
|
| 1754 |
|
|
@@ -1770,7 +1555,7 @@ async def health_check():
|
|
| 1770 |
'pymatting_available': PYMATTING_AVAILABLE,
|
| 1771 |
'image_formats': ['png', 'jpg', 'jpeg', 'webp', 'bmp'],
|
| 1772 |
'bulk_limit': 20,
|
| 1773 |
-
'version': '
|
| 1774 |
'default_model': get_default_model(),
|
| 1775 |
'default_video_model': DEFAULT_VIDEO_MODEL,
|
| 1776 |
'preload_models': [n for n, _ in PRELOAD_MODELS],
|
|
@@ -1778,7 +1563,7 @@ async def health_check():
|
|
| 1778 |
'matting_models': sorted(MATTING_MODELS),
|
| 1779 |
'pad_models': sorted(PAD_MULTIPLE_32_MODELS),
|
| 1780 |
'frame_batch_endpoint': '/api/process/frames-batch',
|
| 1781 |
-
'editor_endpoints': ['auto-crop', 'inpaint'
|
| 1782 |
}
|
| 1783 |
|
| 1784 |
|
|
@@ -1798,10 +1583,10 @@ async def gen_exc(request, exc):
|
|
| 1798 |
@app.on_event("startup")
|
| 1799 |
async def startup():
|
| 1800 |
logger.info("=" * 60)
|
| 1801 |
-
logger.info("
|
| 1802 |
-
logger.info("
|
| 1803 |
-
logger.info(f"
|
| 1804 |
-
logger.info("
|
| 1805 |
logger.info("=" * 60)
|
| 1806 |
|
| 1807 |
|
|
|
|
| 1 |
#!/usr/bin/env python3
|
| 2 |
+
# app.py - Background Remover v32.0.0
|
| 3 |
+
# Tuned for Hugging Face free tier + Vercel frontend.
|
| 4 |
|
| 5 |
import io
|
| 6 |
import time
|
|
|
|
| 15 |
import shutil
|
| 16 |
import zipfile
|
| 17 |
import json
|
|
|
|
| 18 |
import resource
|
| 19 |
import uuid
|
| 20 |
import base64
|
| 21 |
import urllib.parse
|
|
|
|
| 22 |
from collections import defaultdict, Counter, deque
|
| 23 |
from datetime import datetime
|
| 24 |
from pathlib import Path
|
| 25 |
+
from typing import Dict, Optional, List, Any
|
| 26 |
from concurrent.futures import ThreadPoolExecutor
|
| 27 |
|
| 28 |
import psutil
|
| 29 |
|
| 30 |
# ============================================================
|
| 31 |
+
# ONNX RUNTIME — arena ON, threads limited
|
| 32 |
# ============================================================
|
|
|
|
|
|
|
|
|
|
|
|
|
| 33 |
try:
|
| 34 |
import onnxruntime as _ort
|
| 35 |
_orig_session_init = _ort.InferenceSession.__init__
|
|
|
|
| 40 |
so = _ort.SessionOptions()
|
| 41 |
so.intra_op_num_threads = 2
|
| 42 |
so.inter_op_num_threads = 1
|
| 43 |
+
so.enable_cpu_mem_arena = True
|
| 44 |
+
so.enable_mem_pattern = True
|
| 45 |
so.graph_optimization_level = _ort.GraphOptimizationLevel.ORT_ENABLE_ALL
|
| 46 |
kwargs['sess_options'] = so
|
| 47 |
except Exception:
|
|
|
|
| 49 |
return _orig_session_init(self, *args, **kwargs)
|
| 50 |
|
| 51 |
_ort.InferenceSession.__init__ = _patched_session_init
|
| 52 |
+
print("ONNX Runtime patched: arena=ON, threads=2/1")
|
| 53 |
except Exception as e:
|
| 54 |
+
print(f"Could not patch ONNX Runtime: {e}")
|
| 55 |
|
| 56 |
import numpy as np
|
| 57 |
+
from fastapi import FastAPI, File, UploadFile, HTTPException, Query, Form, Request
|
| 58 |
+
from fastapi.responses import StreamingResponse, JSONResponse, HTMLResponse
|
| 59 |
from fastapi.middleware.cors import CORSMiddleware
|
| 60 |
from PIL import Image, ImageDraw
|
| 61 |
import cv2
|
|
|
|
| 73 |
|
| 74 |
|
| 75 |
# ============================================================
|
| 76 |
+
# MEMORY + ENV
|
| 77 |
# ============================================================
|
| 78 |
class ResolutionManager:
|
| 79 |
RECOMMENDED = {
|
|
|
|
| 135 |
|
| 136 |
|
| 137 |
class ModelManager:
|
| 138 |
+
def __init__(self, max_models: int = 3, memory_limit_mb: int = 8000):
|
| 139 |
self.models: Dict[str, Any] = {}
|
| 140 |
self.load_times: Dict[str, float] = {}
|
| 141 |
self.use_counts: Dict[str, int] = {}
|
|
|
|
| 149 |
try:
|
| 150 |
mem = get_container_memory()
|
| 151 |
if 'cgroup' in mem.get('source', '') and mem['percent'] > 82:
|
| 152 |
+
logger.warning(f"Container RAM {mem['percent']:.1f}% — evicting")
|
| 153 |
self._evict_models()
|
| 154 |
gc.collect()
|
| 155 |
except Exception:
|
|
|
|
| 165 |
gc.collect()
|
| 166 |
try:
|
| 167 |
mem = get_container_memory()
|
| 168 |
+
logger.info(f"After {name}: {mem['used_mb']}MB "
|
| 169 |
f"({mem['percent']:.1f}% of {mem['total_mb']}MB)")
|
| 170 |
except Exception:
|
| 171 |
pass
|
|
|
|
| 181 |
if len(self.models) <= 1:
|
| 182 |
return
|
| 183 |
to_remove = min(self.models.keys(), key=lambda x: self.use_counts.get(x, 0))
|
| 184 |
+
logger.info(f"Evicting LRU: {to_remove}")
|
| 185 |
del self.models[to_remove]
|
| 186 |
self.use_counts.pop(to_remove, None)
|
| 187 |
gc.collect()
|
|
|
|
| 415 |
return 'complex'
|
| 416 |
|
| 417 |
|
|
|
|
|
|
|
|
|
|
| 418 |
def limit_memory():
|
| 419 |
try:
|
| 420 |
soft, hard = resource.getrlimit(resource.RLIMIT_AS)
|
| 421 |
resource.setrlimit(resource.RLIMIT_AS, (12 * 1024 * 1024 * 1024, hard))
|
| 422 |
+
logger.info("Memory limit set to 12GB")
|
| 423 |
except Exception as e:
|
| 424 |
+
logger.warning(f"Could not set memory limit: {e}")
|
| 425 |
|
| 426 |
|
| 427 |
limit_memory()
|
|
|
|
| 439 |
cpu_executor = ThreadPoolExecutor(max_workers=2)
|
| 440 |
|
| 441 |
|
|
|
|
|
|
|
|
|
|
| 442 |
content_classifier = ContentClassifier()
|
| 443 |
resolution_manager = ResolutionManager()
|
| 444 |
+
model_manager = ModelManager(max_models=3, memory_limit_mb=8000)
|
| 445 |
mask_checker = MaskQualityChecker()
|
| 446 |
hair_refiner = HairRefiner()
|
| 447 |
graphic_processor = GraphicProcessor()
|
|
|
|
| 450 |
# ============================================================
|
| 451 |
# MODEL CONFIG
|
| 452 |
# ============================================================
|
| 453 |
+
# Only preload the two most useful models. isnet lazy-loads.
|
| 454 |
PRELOAD_MODELS = [
|
| 455 |
("u2netp", "~4.7 MB disk / ~180 MB RAM"),
|
|
|
|
| 456 |
("birefnet-general-lite", "~12 MB disk / ~250 MB RAM"),
|
| 457 |
]
|
| 458 |
|
| 459 |
LAZY_MODELS = {
|
| 460 |
+
"isnet-general-use": "~176 MB disk / ~800 MB RAM",
|
| 461 |
"silueta": "~43 MB disk / ~200 MB RAM",
|
| 462 |
"u2net": "~176 MB disk / ~800 MB RAM",
|
| 463 |
"u2net_human_seg": "~176 MB disk / ~800 MB RAM",
|
|
|
|
| 517 |
|
| 518 |
CLASSICAL_PROCESSORS = {'flood_fill', 'graphic_classical'}
|
| 519 |
DEFAULT_MODEL_NAME = 'isnet-general-use'
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 520 |
DEFAULT_VIDEO_MODEL = 'birefnet-general-lite'
|
| 521 |
|
| 522 |
DISPLAY_NAMES = {
|
|
|
|
| 527 |
'u2net_human_seg': 'U2-Net Human',
|
| 528 |
'u2net_cloth_seg': 'U2-Net Cloth',
|
| 529 |
'silueta': 'Silueta',
|
| 530 |
+
'isnet-general-use': 'ISNet',
|
| 531 |
'isnet-anime': 'ISNet Anime',
|
| 532 |
'birefnet-general': 'BiRefNet',
|
| 533 |
'birefnet-general-lite': 'BiRefNet Lite',
|
|
|
|
| 535 |
}
|
| 536 |
|
| 537 |
|
|
|
|
|
|
|
|
|
|
| 538 |
def _load_session(name: str):
|
| 539 |
from rembg import new_session
|
| 540 |
resolved = MODEL_ALIASES.get(name, name)
|
| 541 |
if resolved not in VALID_REMBG_MODELS:
|
| 542 |
+
logger.warning(f"'{resolved}' not valid -> isnet-general-use")
|
| 543 |
resolved = 'isnet-general-use'
|
| 544 |
if resolved != name:
|
| 545 |
+
logger.info(f"Alias: {name} -> {resolved}")
|
| 546 |
+
logger.info(f"Loading {resolved}...")
|
| 547 |
return new_session(resolved)
|
| 548 |
|
| 549 |
|
|
|
|
| 664 |
return None, diagnostics
|
| 665 |
|
| 666 |
|
|
|
|
|
|
|
|
|
|
| 667 |
def process_graphic_classical(image):
|
| 668 |
try:
|
| 669 |
if image.mode != 'RGBA':
|
|
|
|
| 735 |
PYMATTING_AVAILABLE = False
|
| 736 |
try:
|
| 737 |
from pymatting import cutout, estimate_foreground_ml, estimate_alpha_cf
|
| 738 |
+
from pymatting.util.util import stack_images
|
| 739 |
PYMATTING_AVAILABLE = True
|
| 740 |
+
logger.info("pymatting available")
|
| 741 |
except ImportError as e:
|
| 742 |
+
logger.warning(f"pymatting not available: {e}")
|
| 743 |
|
| 744 |
|
| 745 |
def pre_download_models():
|
| 746 |
+
logger.info(f"Preloading {len(PRELOAD_MODELS)} hot models...")
|
| 747 |
for name, size in PRELOAD_MODELS:
|
| 748 |
try:
|
| 749 |
logger.info(f" Downloading {name} ({size})...")
|
| 750 |
session = model_manager.get_model(name, _load_session)
|
| 751 |
if session is not None:
|
| 752 |
+
logger.info(f" {name} ready")
|
| 753 |
except Exception as e:
|
| 754 |
+
logger.warning(f" {name}: {e}")
|
| 755 |
+
logger.info("Preload complete")
|
| 756 |
|
| 757 |
|
| 758 |
def initialize_models():
|
|
|
|
| 773 |
default_model_name = loaded[0] if loaded else 'u2netp'
|
| 774 |
AI_AVAILABLE = len(loaded) > 0
|
| 775 |
if AI_AVAILABLE:
|
| 776 |
+
logger.info(f"AI ready with {len(loaded)} models")
|
| 777 |
except Exception as e:
|
| 778 |
logger.error(f"AI setup failed: {e}")
|
| 779 |
AI_AVAILABLE = False
|
|
|
|
| 783 |
global remove_func
|
| 784 |
if remove_func is None:
|
| 785 |
return
|
| 786 |
+
logger.info("Pre-warming sessions...")
|
| 787 |
try:
|
| 788 |
test_img = Image.new('RGB', (64, 64), (255, 255, 255))
|
| 789 |
ImageDraw.Draw(test_img).rectangle([20, 20, 44, 44], fill=(220, 40, 40))
|
|
|
|
| 793 |
continue
|
| 794 |
try:
|
| 795 |
t0 = time.time()
|
|
|
|
|
|
|
| 796 |
result = remove_func(test_img, session=session, alpha_matting=False)
|
| 797 |
+
dt = time.time() - t0
|
| 798 |
if result.mode != 'RGBA':
|
| 799 |
result = result.convert('RGBA')
|
| 800 |
arr = np.array(result)
|
| 801 |
cy, cx = arr.shape[0] // 2, arr.shape[1] // 2
|
| 802 |
ca = int(arr[cy, cx, 3]); corner = int(arr[2, 2, 3])
|
| 803 |
if ca > 200 and corner < 50:
|
| 804 |
+
logger.info(f" {name} OK ({dt:.2f}s)")
|
|
|
|
| 805 |
else:
|
| 806 |
+
logger.warning(f" {name} WRONG (center={ca} corner={corner})")
|
| 807 |
except Exception as e:
|
| 808 |
+
logger.warning(f" {name}: {type(e).__name__}: {e}")
|
| 809 |
except Exception as e:
|
| 810 |
logger.warning(f"Pre-warm skipped: {e}")
|
| 811 |
|
| 812 |
|
| 813 |
logger.info("=" * 60)
|
| 814 |
+
logger.info("Background Remover v32.0.0")
|
| 815 |
logger.info("=" * 60)
|
| 816 |
pre_download_models()
|
| 817 |
initialize_models()
|
| 818 |
_prewarm_models()
|
| 819 |
mem = get_container_memory()
|
| 820 |
+
logger.info(f"AI: {AI_AVAILABLE} | Default: {default_model_name}")
|
| 821 |
+
logger.info(f"RAM: {mem['used_mb']}MB / {mem['total_mb']}MB ({mem['percent']:.1f}%)")
|
| 822 |
+
logger.info(f"Video default: {DEFAULT_VIDEO_MODEL}")
|
| 823 |
logger.info("=" * 60)
|
| 824 |
|
| 825 |
|
|
|
|
| 844 |
|
| 845 |
|
| 846 |
# ============================================================
|
| 847 |
+
# ROUTER
|
| 848 |
# ============================================================
|
| 849 |
def detect_hair_complexity(image):
|
| 850 |
+
"""
|
| 851 |
+
FIXED: Previous thresholds fired 'wispy_hair' on almost every image.
|
| 852 |
+
Now requires both a high composite score AND a high fine/coarse ratio.
|
| 853 |
+
Typical photos land in 'simple_hair' or 'complex_hair'.
|
| 854 |
+
"""
|
| 855 |
try:
|
| 856 |
gray = image.convert('L').resize((100, 100), Image.Resampling.LANCZOS)
|
| 857 |
np_img = np.array(gray, dtype=np.float32)
|
|
|
|
| 860 |
fine = np.sum(em > 10) / 10000
|
| 861 |
coarse = np.sum(em > 40) / 10000
|
| 862 |
fr = fine / (coarse + 0.001)
|
| 863 |
+
# Reduced multipliers so the score is not inflated by luminance variance.
|
| 864 |
+
score = min(100, fr * 15 + np.var(np_img) / 8000 * 20)
|
| 865 |
+
|
| 866 |
+
if score > 65 and fr > 0.80:
|
| 867 |
return {'complexity_score': score, 'route': 'wispy_hair'}
|
| 868 |
+
if score > 45 and fr > 0.55:
|
| 869 |
return {'complexity_score': score, 'route': 'complex_hair'}
|
| 870 |
return {'complexity_score': score, 'route': 'simple_hair'}
|
| 871 |
except Exception:
|
|
|
|
| 982 |
# IMAGE PIPELINE
|
| 983 |
# ============================================================
|
| 984 |
def process_image_sync(image_data, model_name=None, transparent=True,
|
| 985 |
+
max_size=1024, quality=95, preserve_hair=True,
|
| 986 |
use_alpha_matting=True, border_pixels=20, sharpen=True):
|
| 987 |
req = str(uuid.uuid4())[:8]
|
| 988 |
logger.info(f"[REQ {req}] === START ===")
|
|
|
|
| 1008 |
'media_type': mt, 'extension': ext,
|
| 1009 |
'original_size': original_size, 'processed_size': image.size,
|
| 1010 |
'fragment_ratio': 0.0, 'quality_score': 100.0,
|
| 1011 |
+
'quality_label': 'Skipped', 'quality_emoji': '',
|
| 1012 |
'has_fragments': False, 'issues': [],
|
| 1013 |
'needs_reprocess': False, 'router_enabled': False,
|
| 1014 |
'router_score': 0.0, 'router_reason': 'already transparent',
|
|
|
|
| 1084 |
frag = detect_fragments(result)
|
| 1085 |
|
| 1086 |
qs = frag['quality_score']
|
| 1087 |
+
if qs >= 90: ql = 'Excellent'
|
| 1088 |
+
elif qs >= 75: ql = 'Good'
|
| 1089 |
+
elif qs >= 50: ql = 'Fair'
|
| 1090 |
+
else: ql = 'Poor'
|
| 1091 |
|
| 1092 |
buf = io.BytesIO()
|
| 1093 |
if transparent:
|
|
|
|
| 1106 |
'original_size': original_size, 'processed_size': result.size,
|
| 1107 |
'fragment_ratio': frag['fragment_ratio'],
|
| 1108 |
'quality_score': frag['quality_score'],
|
| 1109 |
+
'quality_label': ql, 'quality_emoji': '',
|
| 1110 |
'has_fragments': frag['has_fragments'], 'issues': frag['issues'],
|
| 1111 |
'needs_reprocess': frag['has_fragments'],
|
| 1112 |
'router_enabled': router_enabled,
|
|
|
|
| 1125 |
raise
|
| 1126 |
|
| 1127 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1128 |
# ============================================================
|
| 1129 |
# FASTAPI
|
| 1130 |
# ============================================================
|
| 1131 |
+
app = FastAPI(title="Background Remover", version="32.0.0")
|
| 1132 |
app.add_middleware(CORSMiddleware, allow_origins=["*"], allow_credentials=True,
|
| 1133 |
allow_methods=["*"], allow_headers=["*"])
|
| 1134 |
|
|
|
|
| 1147 |
@app.middleware("http")
|
| 1148 |
async def rate_limit_middleware(request, call_next):
|
| 1149 |
if request.url.path in ("/api/health", "/", "/docs", "/redoc", "/openapi.json",
|
| 1150 |
+
"/api/warmup", "/api/debug/env", "/api/debug/model-report"):
|
| 1151 |
return await call_next(request)
|
| 1152 |
cid = _get_client_ip(request)
|
| 1153 |
now = datetime.now()
|
|
|
|
| 1170 |
return deco
|
| 1171 |
|
| 1172 |
|
|
|
|
|
|
|
|
|
|
| 1173 |
@app.post("/api/process/image")
|
| 1174 |
@timeout_handler(seconds=90)
|
| 1175 |
async def process_image_api(
|
| 1176 |
file: UploadFile = File(...),
|
| 1177 |
model: str = Query(None),
|
| 1178 |
transparent: bool = Query(True),
|
| 1179 |
+
max_size: int = Query(1024, ge=512, le=2048),
|
| 1180 |
quality: int = Query(95, ge=70, le=100),
|
| 1181 |
preserve_hair: bool = Query(True),
|
| 1182 |
use_alpha_matting: bool = Query(True),
|
|
|
|
| 1209 |
|
| 1210 |
pt = time.time() - t0
|
| 1211 |
fn = Path(file.filename).stem
|
|
|
|
| 1212 |
headers = {
|
| 1213 |
'X-Processing-Time': f'{pt:.3f}',
|
| 1214 |
'X-Model-Used': result_data['model'],
|
|
|
|
| 1216 |
'X-Transparent': str(result_data['transparent']).lower(),
|
| 1217 |
'X-Quality-Score': f"{result_data['quality_score']:.1f}",
|
| 1218 |
'X-Quality-Label': result_data.get('quality_label', ''),
|
|
|
|
| 1219 |
'X-Fragment-Ratio': f"{result_data.get('fragment_ratio', 0):.3f}",
|
| 1220 |
'X-Needs-Reprocess': str(result_data.get('needs_reprocess', False)).lower(),
|
| 1221 |
'X-Router-Enabled': str(result_data.get('router_enabled', False)).lower(),
|
|
|
|
| 1234 |
raise HTTPException(500, f"Processing failed: {str(e)[:200]}")
|
| 1235 |
|
| 1236 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1237 |
@app.post("/api/process/bulk")
|
| 1238 |
@timeout_handler(seconds=300)
|
| 1239 |
async def process_bulk_images(
|
|
|
|
| 1303 |
raise HTTPException(500, str(e)[:200])
|
| 1304 |
|
| 1305 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1306 |
# ============================================================
|
| 1307 |
+
# FRAME-BATCH ENDPOINT
|
| 1308 |
# ============================================================
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1309 |
@app.post("/api/process/frames-batch")
|
| 1310 |
async def process_frames_batch(
|
| 1311 |
files: List[UploadFile] = File(...),
|
| 1312 |
model: str = Query(DEFAULT_VIDEO_MODEL),
|
| 1313 |
transparent: bool = Query(True),
|
| 1314 |
):
|
| 1315 |
+
# HARD CAP to prevent OOM on free tier.
|
| 1316 |
+
if len(files) > 4:
|
| 1317 |
+
raise HTTPException(400, "Max 4 frames per batch")
|
| 1318 |
+
|
| 1319 |
resolved = MODEL_ALIASES.get(model, model)
|
| 1320 |
if resolved not in VALID_REMBG_MODELS:
|
| 1321 |
resolved = DEFAULT_VIDEO_MODEL
|
|
|
|
| 1324 |
if session is None:
|
| 1325 |
raise HTTPException(503, f"No session for {resolved}")
|
| 1326 |
|
|
|
|
|
|
|
| 1327 |
fallback_chain = [resolved]
|
| 1328 |
for candidate in ('birefnet-general-lite', 'isnet-general-use', 'u2netp'):
|
| 1329 |
if candidate not in fallback_chain:
|
|
|
|
| 1359 |
min_a = int(alpha.min())
|
| 1360 |
max_a = int(alpha.max())
|
| 1361 |
|
|
|
|
| 1362 |
if max_a < 10:
|
| 1363 |
+
logger.warning(f"[frames-batch] {candidate} blank (max_a={max_a}); "
|
| 1364 |
+
f"frame {idx} trying next")
|
|
|
|
|
|
|
| 1365 |
continue
|
|
|
|
|
|
|
| 1366 |
if min_a > 245:
|
| 1367 |
+
logger.warning(f"[frames-batch] {candidate} no-op (min_a={min_a}); "
|
| 1368 |
+
f"frame {idx} trying next")
|
|
|
|
|
|
|
| 1369 |
continue
|
| 1370 |
|
| 1371 |
result = res
|
|
|
|
| 1373 |
break
|
| 1374 |
|
| 1375 |
if result is None:
|
|
|
|
|
|
|
| 1376 |
logger.error(f"[frames-batch] frame {idx}: all models failed")
|
| 1377 |
buf = io.BytesIO()
|
| 1378 |
img.convert("RGBA").save(buf, format="PNG", compress_level=3)
|
|
|
|
| 1399 |
})
|
| 1400 |
except Exception as e:
|
| 1401 |
logger.error(f"[frames-batch] frame {idx} failed: {e}")
|
| 1402 |
+
results.append({"data": "", "format": "png", "error": str(e)[:120]})
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1403 |
|
| 1404 |
used = next((r.get("model") for r in results if r.get("model")), resolved)
|
| 1405 |
+
return {"frames": results, "count": len(results), "model": used}
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1406 |
|
| 1407 |
|
| 1408 |
# ============================================================
|
|
|
|
| 1481 |
raise HTTPException(500, f"Inpainting failed: {str(e)[:200]}")
|
| 1482 |
|
| 1483 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1484 |
def _hex_to_rgb(hex_str):
|
| 1485 |
h = hex_str.lstrip('#')
|
| 1486 |
if len(h) == 3:
|
|
|
|
| 1489 |
|
| 1490 |
|
| 1491 |
# ============================================================
|
| 1492 |
+
# DEBUG / HEALTH / WARMUP
|
| 1493 |
# ============================================================
|
| 1494 |
+
@app.get("/api/warmup")
|
| 1495 |
+
async def warmup():
|
| 1496 |
+
"""Lightweight keep-alive endpoint for Vercel cron."""
|
| 1497 |
+
return {"status": "warm", "ts": datetime.now().isoformat()}
|
| 1498 |
+
|
| 1499 |
+
|
| 1500 |
@app.get("/api/debug/env")
|
| 1501 |
async def debug_env():
|
| 1502 |
m = get_container_memory()
|
|
|
|
| 1515 |
'matting_models': sorted(MATTING_MODELS)}
|
| 1516 |
|
| 1517 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1518 |
@app.get("/")
|
| 1519 |
async def home():
|
| 1520 |
stats = model_manager.get_stats()
|
| 1521 |
mem = get_container_memory()
|
| 1522 |
return HTMLResponse(f"""
|
| 1523 |
<!DOCTYPE html><html><head><meta charset="UTF-8">
|
| 1524 |
+
<title>Background Remover v32.0.0</title>
|
| 1525 |
+
<style>body{{font-family:system-ui;background:#0A0A0A;
|
| 1526 |
color:#fff;min-height:100vh;display:flex;align-items:center;justify-content:center;
|
| 1527 |
+
margin:0;padding:20px}}.c{{max-width:900px;background:#111;
|
| 1528 |
+
border-radius:20px;padding:40px;text-align:center;
|
| 1529 |
+
border:1px solid #222}}.s{{padding:15px;border-radius:12px;
|
| 1530 |
+
margin:20px 0;background:#1a1a1a}}a{{color:#7ED957}}</style>
|
| 1531 |
</head><body><div class="c">
|
| 1532 |
+
<h1>Background Remover v32.0.0</h1>
|
| 1533 |
+
<div class="s">AI Online · {len(stats['loaded_models'])} models ·
|
| 1534 |
RAM {mem['percent']:.0f}% ({mem['used_mb']}/{mem['total_mb']}MB)</div>
|
| 1535 |
<p>Video default: <strong>{DEFAULT_VIDEO_MODEL}</strong></p>
|
| 1536 |
+
<p><a href="/docs">API Docs</a></p>
|
| 1537 |
</div></body></html>
|
| 1538 |
""")
|
| 1539 |
|
|
|
|
| 1555 |
'pymatting_available': PYMATTING_AVAILABLE,
|
| 1556 |
'image_formats': ['png', 'jpg', 'jpeg', 'webp', 'bmp'],
|
| 1557 |
'bulk_limit': 20,
|
| 1558 |
+
'version': '32.0.0',
|
| 1559 |
'default_model': get_default_model(),
|
| 1560 |
'default_video_model': DEFAULT_VIDEO_MODEL,
|
| 1561 |
'preload_models': [n for n, _ in PRELOAD_MODELS],
|
|
|
|
| 1563 |
'matting_models': sorted(MATTING_MODELS),
|
| 1564 |
'pad_models': sorted(PAD_MULTIPLE_32_MODELS),
|
| 1565 |
'frame_batch_endpoint': '/api/process/frames-batch',
|
| 1566 |
+
'editor_endpoints': ['auto-crop', 'inpaint'],
|
| 1567 |
}
|
| 1568 |
|
| 1569 |
|
|
|
|
| 1583 |
@app.on_event("startup")
|
| 1584 |
async def startup():
|
| 1585 |
logger.info("=" * 60)
|
| 1586 |
+
logger.info("v32.0.0 ready")
|
| 1587 |
+
logger.info("Video pipeline: client-side frame batching")
|
| 1588 |
+
logger.info(f"Video default model: {DEFAULT_VIDEO_MODEL}")
|
| 1589 |
+
logger.info("ONNX: arena=ON, mem_pattern=ON")
|
| 1590 |
logger.info("=" * 60)
|
| 1591 |
|
| 1592 |
|