Refactor code structure for improved readability and maintainability
This commit is contained in:
@@ -0,0 +1,21 @@
|
||||
# RunPod Serverless worker: background removal (rembg / U²-Net), CPU-only.
|
||||
# Slim Python base — NO CUDA — so it never hits the driver/kernel issues the GPU
|
||||
# models did. rembg pulls opencv (needs libgl1/libglib). Build context = repo root.
|
||||
FROM python:3.11-slim
|
||||
|
||||
ENV PYTHONUNBUFFERED=1 \
|
||||
PIP_NO_CACHE_DIR=1
|
||||
|
||||
RUN apt-get update && apt-get install -y --no-install-recommends \
|
||||
libgl1 libglib2.0-0 && \
|
||||
rm -rf /var/lib/apt/lists/*
|
||||
|
||||
WORKDIR /app
|
||||
|
||||
RUN pip install --upgrade pip
|
||||
COPY workers/bg-remove/requirements.txt .
|
||||
RUN pip install -r requirements.txt
|
||||
|
||||
COPY workers/bg-remove/handler.py .
|
||||
|
||||
CMD ["python3", "-u", "handler.py"]
|
||||
@@ -0,0 +1,17 @@
|
||||
# Background-removal worker — rembg / U²-Net (CPU)
|
||||
|
||||
Replaces the unreliable in-browser @imgly remover with a RunPod endpoint. Set its
|
||||
`/runsync` URL as **`RUNPOD_BG_REMOVE_URL`**, and turn it on with
|
||||
**`NEXT_PUBLIC_APG_RUNPOD_BG=true`** (without the flag the app keeps using the
|
||||
in-browser fallback).
|
||||
|
||||
## Contract
|
||||
`{"input": {"image": "<base64>", "model_name": "u2net"?}}` → `{"image": "<base64 PNG, transparent background>"}`.
|
||||
|
||||
Models: `u2net` (default, general), `u2netp` (lighter), `u2net_human_seg` (people), `isnet-general-use` (sharper edges).
|
||||
|
||||
## RunPod setup
|
||||
1. New Serverless endpoint → connect this repo → **Dockerfile path** `workers/bg-remove/Dockerfile`.
|
||||
2. Any GPU is fine (it runs on CPU) — pick the cheapest. Max workers ≥ 1.
|
||||
3. Optional volume to cache the ~170 MB model (`U2NET_HOME` → `/runpod-volume/u2net`).
|
||||
4. Copy the `…/runsync` URL → Vercel `RUNPOD_BG_REMOVE_URL`, add `NEXT_PUBLIC_APG_RUNPOD_BG=true`, redeploy.
|
||||
@@ -0,0 +1,66 @@
|
||||
"""
|
||||
RunPod Serverless worker: background removal (U²-Net via rembg).
|
||||
|
||||
Replaces the flaky in-browser @imgly remover. Runs rembg on CPU (U²-Net is tiny,
|
||||
~1-3s) so it needs NO CUDA — it can't hit the driver/kernel problems the GPU
|
||||
models had, and it's very reliable.
|
||||
|
||||
Contract (matches the app's rpRemoveBackground — no app change needed):
|
||||
{"input": {
|
||||
"image": "<base64>", # required (raw base64 or data: URI)
|
||||
"model_name": "u2net" # optional: u2net | u2netp | u2net_human_seg | isnet-general-use
|
||||
}}
|
||||
Returns: {"image": "<base64 PNG, RGBA with the background transparent>"}.
|
||||
"""
|
||||
|
||||
import base64
|
||||
import io
|
||||
import os
|
||||
|
||||
_V = "/runpod-volume"
|
||||
if os.path.isdir(_V):
|
||||
os.environ.setdefault("U2NET_HOME", os.path.join(_V, "u2net"))
|
||||
|
||||
import runpod # noqa: E402
|
||||
from PIL import Image # noqa: E402
|
||||
from rembg import new_session, remove # noqa: E402
|
||||
|
||||
DEFAULT_MODEL = os.environ.get("BG_REMOVE_MODEL", "u2net")
|
||||
_sessions = {}
|
||||
|
||||
|
||||
def _get_session(name):
|
||||
if name not in _sessions:
|
||||
_sessions[name] = new_session(name)
|
||||
return _sessions[name]
|
||||
|
||||
|
||||
def _b64_to_image(data):
|
||||
if "," in data:
|
||||
data = data.split(",", 1)[1]
|
||||
return Image.open(io.BytesIO(base64.b64decode(data))).convert("RGBA")
|
||||
|
||||
|
||||
def _image_to_b64(img):
|
||||
buf = io.BytesIO()
|
||||
img.save(buf, format="PNG")
|
||||
return base64.b64encode(buf.getvalue()).decode("utf-8")
|
||||
|
||||
|
||||
def handler(event):
|
||||
inp = (event or {}).get("input") or {}
|
||||
image_b64 = inp.get("image")
|
||||
if not image_b64:
|
||||
return {"error": "Missing 'image' (base64)."}
|
||||
name = (inp.get("model_name") or DEFAULT_MODEL).strip() or DEFAULT_MODEL
|
||||
try:
|
||||
img = _b64_to_image(image_b64)
|
||||
out = remove(img, session=_get_session(name)) # RGBA, background alpha=0
|
||||
if out.mode != "RGBA":
|
||||
out = out.convert("RGBA")
|
||||
return {"image": _image_to_b64(out)}
|
||||
except Exception as exc:
|
||||
return {"error": f"{type(exc).__name__}: {exc}"}
|
||||
|
||||
|
||||
runpod.serverless.start({"handler": handler})
|
||||
@@ -0,0 +1,6 @@
|
||||
# Background-removal worker (CPU). rembg pulls onnxruntime (CPU) + opencv itself.
|
||||
rembg==2.0.59
|
||||
onnxruntime==1.17.3
|
||||
numpy<2
|
||||
Pillow==10.4.0
|
||||
runpod==1.7.9
|
||||
@@ -0,0 +1 @@
|
||||
{"input": {"image": "iVBORw0KGgoAAAANSUhEUgAAAQAAAAEACAIAAADTED8xAAACAElEQVR42u3TQQ0AAAjEsJONCEQglTcaaFIFS5bqgbciAQYAA4ABwABgADAAGAAMAAYAA4ABwABgADAAGAAMAAYAA4ABwABgADAAGAAMAAYAA4ABwABgADAAGAAMAAYAA4ABwABgADAAGAAMAAYAA4ABwABgADAAGAAMAAYAA4ABMIAKGAAMAAYAA4ABwABgADAAGAAMAAYAA4ABwABgADAAGAAMAAYAA4ABwABgADAAGAAMAAYAA4ABwABgADAAGAAMAAYAA4ABwABgADAAGAAMAAYAA4ABwABgADAAGAAMAAbAAGAAMAAYAAwABgADgAHAAGAAMAAYAAwABgADgAHAAGAAMAAYAAwABgADgAHAAGAAMAAYAAwABgADgAHAAGAAMAAYAAwABgADgAHAAGAAMAAYAAwABgADgAHAAGAAMAAGUAEDgAHAAGAAMAAYAAwABgADgAHAAGAAMAAYAAwABgADgAHAAGAAMAAYAAwABgADgAHAAGAAMAAYAAwABgADgAHAAGAAMAAYAAwABgADgAHAAGAAMAAYAAwABgADgAHAABgADAAGAAOAAcAAYAAwABgADAAGAAOAAcAAYAAwABgADAAGAAOAAcAAYAAwABgADAAGAAOAAcAAYAAwABgADAAGAAOAAcAAYAAwABgADAAGAAOAAcAAYAAwAFwLd10YnS3DbrYAAAAASUVORK5CYII=", "model_name": "u2net"}}
|
||||
Reference in New Issue
Block a user