Refactor code structure for improved readability and maintainability

This commit is contained in:
2026-07-22 22:58:28 +05:30
parent d1308bf149
commit bc8bf2007d
99 changed files with 4784 additions and 111 deletions
+23
View File
@@ -0,0 +1,23 @@
FROM python:3.11-slim
WORKDIR /app
# ffmpeg (with arnndn/afftdn filters) + curl for model download
RUN apt-get update && \
apt-get install -y --no-install-recommends ffmpeg ca-certificates curl && \
rm -rf /var/lib/apt/lists/*
# Download a standard RNNoise model for ffmpeg's arnndn filter.
# Source: GregorR/rnnoise-models (beguiling-drafts, a widely used general model).
# Non-fatal: if the download fails, handler.py falls back to afftdn.
RUN mkdir -p /app/models && \
( curl -fsSL -o /app/models/rnnoise.rnnn \
https://raw.githubusercontent.com/GregorR/rnnoise-models/master/beguiling-drafts-2018-08-30/bd.rnnn \
|| echo "WARNING: RNNoise model download failed; will fall back to afftdn" )
COPY workers/audio-denoise/requirements.txt .
RUN pip install --no-cache-dir -r requirements.txt
COPY workers/audio-denoise/handler.py .
CMD ["python3", "-u", "handler.py"]
+26
View File
@@ -0,0 +1,26 @@
# Audio Denoise Worker (RNNoise / arnndn)
RunPod Serverless CPU worker that removes background noise from audio using
ffmpeg's RNNoise filter (`arnndn`), with a built-in FFT denoiser (`afftdn`)
as an automatic fallback.
## Contract
- Input: `{"audio": "<base64>"}` — any format/sample rate. A `data:...;base64,`
prefix is stripped automatically.
- Output: `{"audio": "<base64 WAV>"}` — denoised, 48 kHz mono WAV.
- Error: `{"error": "<message>"}`.
## Pipeline
decode base64 -> temp input file -> `ffmpeg -i in -af arnndn=m=/app/models/rnnoise.rnnn -ar 48000 -ac 1 out.wav` -> read -> base64.
## Deploy
- Dockerfile path: `/workers/audio-denoise/Dockerfile`
- Build context: repo root (`docker build -f workers/audio-denoise/Dockerfile .`)
- GPU: none (CPU-only worker)
- App env var to set: `RUNPOD_AUDIO_DENOISE_URL` (your deployed endpoint URL)
## Model
The Dockerfile downloads the `beguiling-drafts` RNNoise model from
`GregorR/rnnoise-models` to `/app/models/rnnoise.rnnn`. If that download
fails at build time, the handler transparently falls back to ffmpeg's
built-in `afftdn` FFT denoiser (no model file needed).
+70
View File
@@ -0,0 +1,70 @@
_V = "/runpod-volume"
import os
if os.path.isdir(_V):
os.environ.setdefault("HF_HOME", os.path.join(_V, "huggingface"))
import base64
import subprocess
import tempfile
import runpod
# Path where the Dockerfile places the RNNoise model file.
RNNN_MODEL = "/app/models/rnnoise.rnnn"
def _decode_b64(s):
if isinstance(s, str) and "," in s and s.strip().lower().startswith("data:"):
s = s.split(",", 1)[1]
return base64.b64decode(s)
def _denoise(in_path, out_path):
"""Run ffmpeg. Prefer RNNoise (arnndn) when a model file is present,
otherwise fall back to the built-in FFT denoiser (afftdn)."""
if os.path.isfile(RNNN_MODEL):
af = "arnndn=m=" + RNNN_MODEL
else:
af = "afftdn=nf=-25"
cmd = [
"ffmpeg", "-y", "-hide_banner", "-nostdin",
"-i", in_path,
"-af", af,
"-ar", "48000", "-ac", "1",
"-f", "wav",
out_path,
]
proc = subprocess.run(cmd, capture_output=True)
if proc.returncode != 0:
raise RuntimeError(
"ffmpeg failed: " + proc.stderr.decode("utf-8", "replace")[-800:]
)
def handler(event):
try:
inp = (event or {}).get("input") or {}
audio_b64 = inp.get("audio")
if not audio_b64:
return {"error": "Missing 'audio' in input"}
raw = _decode_b64(audio_b64)
with tempfile.TemporaryDirectory() as td:
in_path = os.path.join(td, "in")
out_path = os.path.join(td, "out.wav")
with open(in_path, "wb") as f:
f.write(raw)
_denoise(in_path, out_path)
with open(out_path, "rb") as f:
out_bytes = f.read()
return {"audio": base64.b64encode(out_bytes).decode("ascii")}
except Exception as exc:
return {"error": f"{type(exc).__name__}: {exc}"}
runpod.serverless.start({"handler": handler})
+1
View File
@@ -0,0 +1 @@
runpod==1.7.9
File diff suppressed because one or more lines are too long
+21
View File
@@ -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"]
+17
View File
@@ -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.
+66
View File
@@ -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})
+6
View File
@@ -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
+1
View File
@@ -0,0 +1 @@
{"input": {"image": "iVBORw0KGgoAAAANSUhEUgAAAQAAAAEACAIAAADTED8xAAACAElEQVR42u3TQQ0AAAjEsJONCEQglTcaaFIFS5bqgbciAQYAA4ABwABgADAAGAAMAAYAA4ABwABgADAAGAAMAAYAA4ABwABgADAAGAAMAAYAA4ABwABgADAAGAAMAAYAA4ABwABgADAAGAAMAAYAA4ABwABgADAAGAAMAAYAA4ABMIAKGAAMAAYAA4ABwABgADAAGAAMAAYAA4ABwABgADAAGAAMAAYAA4ABwABgADAAGAAMAAYAA4ABwABgADAAGAAMAAYAA4ABwABgADAAGAAMAAYAA4ABwABgADAAGAAMAAbAAGAAMAAYAAwABgADgAHAAGAAMAAYAAwABgADgAHAAGAAMAAYAAwABgADgAHAAGAAMAAYAAwABgADgAHAAGAAMAAYAAwABgADgAHAAGAAMAAYAAwABgADgAHAAGAAMAAGUAEDgAHAAGAAMAAYAAwABgADgAHAAGAAMAAYAAwABgADgAHAAGAAMAAYAAwABgADgAHAAGAAMAAYAAwABgADgAHAAGAAMAAYAAwABgADgAHAAGAAMAAYAAwABgADgAHAABgADAAGAAOAAcAAYAAwABgADAAGAAOAAcAAYAAwABgADAAGAAOAAcAAYAAwABgADAAGAAOAAcAAYAAwABgADAAGAAOAAcAAYAAwABgADAAGAAOAAcAAYAAwAFwLd10YnS3DbrYAAAAASUVORK5CYII=", "model_name": "u2net"}}
+24
View File
@@ -0,0 +1,24 @@
FROM nvidia/cuda:12.1.1-cudnn8-runtime-ubuntu22.04
ENV DEBIAN_FRONTEND=noninteractive \
PIP_NO_CACHE_DIR=1 \
PYTHONUNBUFFERED=1
RUN apt-get update && apt-get install -y --no-install-recommends \
python3 python3-pip python3-dev \
libgl1 libglib2.0-0 \
git ca-certificates \
&& rm -rf /var/lib/apt/lists/* \
&& ln -sf /usr/bin/python3 /usr/bin/python
WORKDIR /app
COPY workers/colorize/requirements.txt .
RUN pip3 install --upgrade pip && pip3 install -r requirements.txt
# Install runpod LAST so modelscope's dependency pins can't break `import runpod`
# at worker startup (the cause of "worker never becomes ready, no logs").
RUN pip3 install runpod==1.7.9
COPY workers/colorize/handler.py .
CMD ["python3", "-u", "handler.py"]
+33
View File
@@ -0,0 +1,33 @@
# DDColor Colorization Worker (#8)
RunPod Serverless worker that colorizes black-and-white images using the
ModelScope DDColor pipeline (`damo/cv_ddcolor_image-colorization`).
## Contract
Input:
```json
{ "image": "<base64 (optionally data:...;base64, prefixed)>", "input_size": 512 }
```
Output:
```json
{ "image": "<base64 PNG, same dimensions as input>" }
```
On error: `{ "error": "..." }`.
## Deploy
- Dockerfile path: `/workers/colorize/Dockerfile`
- Build context: repo root (`docker build -f workers/colorize/Dockerfile .`)
- GPU: medium (RTX 3080 / 3090, ~10-12 GB VRAM)
- App env var to set: `RUNPOD_COLORIZE_URL` (point your app at this endpoint)
## Notes
- Model weights are cached to the network volume via `MODELSCOPE_CACHE`
(`/runpod-volume/modelscope`) when `/runpod-volume` is mounted, so the first
request downloads and subsequent cold starts reuse the cached weights.
- `input_size` controls the DDColor internal working resolution (default 512).
The output is always resized back to the exact input dimensions.
+80
View File
@@ -0,0 +1,80 @@
_V = "/runpod-volume"
import os
if os.path.isdir(_V):
os.environ.setdefault("HF_HOME", os.path.join(_V, "huggingface"))
os.environ.setdefault("MODELSCOPE_CACHE", os.path.join(_V, "modelscope"))
import base64
import runpod
_pipeline = None
def _get_pipeline():
global _pipeline
if _pipeline is None:
from modelscope.pipelines import pipeline
from modelscope.utils.constant import Tasks
_pipeline = pipeline(
Tasks.image_colorization,
model="damo/cv_ddcolor_image-colorization",
)
return _pipeline
def _decode_image(b64):
if "," in b64:
b64 = b64.split(",", 1)[1]
return base64.b64decode(b64)
def handler(event):
try:
import numpy as np
import cv2
inp = (event or {}).get("input") or {}
img_b64 = inp.get("image")
if not img_b64:
return {"error": "ValueError: missing 'image' in input"}
try:
input_size = int(inp.get("input_size", 512))
except (TypeError, ValueError):
input_size = 512
raw = _decode_image(img_b64)
arr = np.frombuffer(raw, dtype=np.uint8)
# Decode to BGR (modelscope DDColor expects BGR numpy array).
bgr = cv2.imdecode(arr, cv2.IMREAD_COLOR)
if bgr is None:
return {"error": "ValueError: could not decode input image"}
orig_h, orig_w = bgr.shape[:2]
pipe = _get_pipeline()
from modelscope.outputs import OutputKeys
result = pipe(bgr, input_size=input_size)
out_bgr = result[OutputKeys.OUTPUT_IMG]
out_bgr = np.asarray(out_bgr)
# Ensure output matches the input dimensions exactly.
if out_bgr.shape[0] != orig_h or out_bgr.shape[1] != orig_w:
out_bgr = cv2.resize(
out_bgr, (orig_w, orig_h), interpolation=cv2.INTER_LANCZOS4
)
# BGR -> RGB, then encode PNG. cv2.imencode expects BGR, so re-convert.
out_rgb = cv2.cvtColor(out_bgr, cv2.COLOR_BGR2RGB)
ok, buf = cv2.imencode(".png", cv2.cvtColor(out_rgb, cv2.COLOR_RGB2BGR))
if not ok:
return {"error": "RuntimeError: PNG encoding failed"}
out_b64 = base64.b64encode(buf.tobytes()).decode("utf-8")
return {"image": out_b64}
except Exception as exc:
return {"error": f"{type(exc).__name__}: {exc}"}
runpod.serverless.start({"handler": handler})
+9
View File
@@ -0,0 +1,9 @@
# NOTE: runpod is installed LAST in the Dockerfile (after modelscope), so
# modelscope's aggressive dependency pins can't break `import runpod` at startup.
torch==2.1.2
torchvision==0.16.2
modelscope==1.9.5
opencv-python-headless==4.9.0.80
pillow==10.2.0
numpy==1.24.4
timm==0.9.12
+1
View File
@@ -0,0 +1 @@
{"input": {"image": "iVBORw0KGgoAAAANSUhEUgAAAgAAAAIACAIAAAB7GkOtAAAFp0lEQVR42u3VMREAMAgAMfRVRO2woK1GagIJCCB3UfDLR74PwEIhAYABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAIABqABgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgBgACoAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGACAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGACAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAwDiAWweAhQwAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAwAAkADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADADAAFQAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAMQAUAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADADAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADADAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAYNQ/e2eZ8ltazgAAAABJRU5ErkJggg=="}}
+26
View File
@@ -0,0 +1,26 @@
# Combined CPU RunPod worker: background removal (rembg) + audio denoise (ffmpeg).
# Slim base, NO CUDA. 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 \
ffmpeg libgl1 libglib2.0-0 curl && \
rm -rf /var/lib/apt/lists/*
WORKDIR /app
# RNNoise model for ffmpeg's arnndn (falls back to afftdn if the download fails).
RUN mkdir -p /app/models && \
(curl -fsSL -o /app/models/rnnoise.rnnn \
https://raw.githubusercontent.com/GregorR/rnnoise-models/master/beguiling-drafts-2018-08-30/bd.rnnn || \
echo "rnnoise model download skipped; afftdn fallback will be used")
RUN pip install --upgrade pip
COPY workers/cpu-tasks/requirements.txt .
RUN pip install -r requirements.txt
COPY workers/cpu-tasks/handler.py .
CMD ["python3", "-u", "handler.py"]
+82
View File
@@ -0,0 +1,82 @@
"""
Combined CPU RunPod worker (no GPU, no torch). Routes on input.task:
"remove-bg" -> rembg (U^2-Net) background removal -> {"image": <png b64>}
"denoise" -> ffmpeg RNNoise/afftdn audio denoise -> {"audio": <wav b64>}
One request runs exactly one model (isolated). If `task` is absent it's inferred
from the payload (audio -> denoise, image -> remove-bg) for backward-compat.
"""
import base64
import io
import os
import subprocess
import tempfile
import runpod
RNNN_MODEL = "/app/models/rnnoise.rnnn"
_sessions = {}
def _decode_b64(s):
if isinstance(s, str) and "," in s:
s = s.split(",", 1)[1]
return base64.b64decode(s)
def _remove_bg(inp):
from PIL import Image
from rembg import new_session, remove
image_b64 = inp.get("image")
if not image_b64:
return {"error": "Missing 'image' (base64)."}
name = (inp.get("model_name") or os.environ.get("BG_REMOVE_MODEL", "u2net")).strip() or "u2net"
if name not in _sessions:
_sessions[name] = new_session(name)
img = Image.open(io.BytesIO(_decode_b64(image_b64))).convert("RGBA")
out = remove(img, session=_sessions[name])
if out.mode != "RGBA":
out = out.convert("RGBA")
buf = io.BytesIO()
out.save(buf, format="PNG")
return {"image": base64.b64encode(buf.getvalue()).decode("utf-8")}
def _denoise(inp):
audio_b64 = inp.get("audio")
if not audio_b64:
return {"error": "Missing 'audio' (base64)."}
raw = _decode_b64(audio_b64)
af = ("arnndn=m=" + RNNN_MODEL) if os.path.isfile(RNNN_MODEL) else "afftdn=nf=-25"
with tempfile.TemporaryDirectory() as td:
in_path = os.path.join(td, "in")
out_path = os.path.join(td, "out.wav")
with open(in_path, "wb") as f:
f.write(raw)
cmd = ["ffmpeg", "-y", "-hide_banner", "-nostdin", "-i", in_path,
"-af", af, "-ar", "48000", "-ac", "1", "-f", "wav", out_path]
proc = subprocess.run(cmd, capture_output=True)
if proc.returncode != 0:
raise RuntimeError("ffmpeg failed: " + proc.stderr.decode("utf-8", "replace")[-600:])
with open(out_path, "rb") as f:
out_bytes = f.read()
return {"audio": base64.b64encode(out_bytes).decode("ascii")}
def handler(event):
inp = (event or {}).get("input") or {}
task = (inp.get("task") or "").strip().lower()
if not task:
task = "denoise" if inp.get("audio") else "remove-bg"
try:
if task == "remove-bg":
return _remove_bg(inp)
if task == "denoise":
return _denoise(inp)
return {"error": f"Unknown task '{task}' for cpu endpoint."}
except Exception as exc:
return {"error": f"{type(exc).__name__}: {exc}"}
runpod.serverless.start({"handler": handler})
+6
View File
@@ -0,0 +1,6 @@
# Combined CPU worker: rembg (bg removal) + ffmpeg denoise. No torch/GPU.
runpod==1.7.9
rembg==2.0.59
onnxruntime==1.17.3
numpy<2
Pillow==10.4.0
+22
View File
@@ -0,0 +1,22 @@
FROM nvidia/cuda:12.1.1-cudnn8-runtime-ubuntu22.04
ENV DEBIAN_FRONTEND=noninteractive
RUN apt-get update && apt-get install -y --no-install-recommends \
python3 python3-pip python3-dev \
libgl1 libglib2.0-0 \
&& rm -rf /var/lib/apt/lists/*
RUN ln -sf /usr/bin/python3 /usr/bin/python
WORKDIR /app
COPY workers/detect/requirements.txt .
RUN pip3 install --no-cache-dir -r requirements.txt
# Bake in the trained construction-material model and point the worker at it.
COPY workers/detect/best.pt .
ENV YOLO_MODEL=/app/best.pt
COPY workers/detect/handler.py .
CMD ["python3", "-u", "handler.py"]
+37
View File
@@ -0,0 +1,37 @@
# Construction-Material Detection Worker (YOLOv8)
RunPod Serverless worker running Ultralytics YOLOv8 object detection for a
custom-trained construction-material classifier.
## Build & Deploy
- Dockerfile path: `/workers/detect/Dockerfile`
- Build context: repo ROOT (the COPY paths are prefixed with `workers/detect/`).
- GPU: light/medium — NVIDIA T4 or A4000 is plenty.
## Model
The trained construction-material model (`best.pt`, 43 MB) is **baked into the
image** (`COPY workers/detect/best.pt``ENV YOLO_MODEL=/app/best.pt`), so the
worker detects your materials out of the box — no volume upload needed.
- `YOLO_MODEL` — override only if you want to swap models (default `/app/best.pt`).
To update the model later, replace `workers/detect/best.pt` and push (RunPod rebuilds).
## App-side env var
Point your application at the deployed endpoint with:
- `RUNPOD_YOLO_URL` = your RunPod serverless endpoint URL.
## Contract
Input:
```json
{ "input": { "image": "<base64>" } }
```
The base64 may include a `data:image/...;base64,` prefix; it is stripped.
Output (pixel coordinates):
```json
{ "detections": [
{ "name": "brick", "confidence": 0.94,
"box": { "x1": 10.0, "y1": 20.0, "x2": 110.0, "y2": 220.0 } }
] }
```
On error: `{ "error": "<Type>: <message>" }`
Binary file not shown.
+64
View File
@@ -0,0 +1,64 @@
_V = "/runpod-volume"
import os
if os.path.isdir(_V):
os.environ.setdefault("HF_HOME", os.path.join(_V, "huggingface"))
import io
import base64
import runpod
from PIL import Image
_model = None
def _load_model():
global _model
if _model is None:
from ultralytics import YOLO
weights = os.environ.get("YOLO_MODEL", "yolov8n.pt")
_model = YOLO(weights)
return _model
def handler(event):
try:
data = (event or {}).get("input") or {}
image_b64 = data.get("image")
if not image_b64:
return {"error": "ValueError: missing 'image' in input"}
if "," in image_b64:
image_b64 = image_b64.split(",", 1)[1]
image_bytes = base64.b64decode(image_b64)
image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
# Custom construction models are often less confident than COCO — use a low
# detection floor (env YOLO_CONF, default 0.15) so real brick/pipe/steel hits
# come through instead of being dropped by ultralytics' default 0.25.
conf = float(data.get("conf") or os.environ.get("YOLO_CONF", "0.15"))
model = _load_model()
results = model.predict(image, conf=conf, verbose=False)
detections = []
for result in results:
boxes = getattr(result, "boxes", None)
if boxes is None:
continue
names = result.names if getattr(result, "names", None) else model.names
for box in boxes:
cls = int(box.cls[0])
conf = float(box.conf[0])
xyxy = box.xyxy[0].tolist()
x1, y1, x2, y2 = [float(v) for v in xyxy]
detections.append({
"name": str(names[cls]),
"confidence": conf,
"box": {"x1": x1, "y1": y1, "x2": x2, "y2": y2},
})
return {"detections": detections}
except Exception as exc:
return {"error": f"{type(exc).__name__}: {exc}"}
runpod.serverless.start({"handler": handler})
+3
View File
@@ -0,0 +1,3 @@
runpod==1.7.9
ultralytics==8.3.40
pillow==10.4.0
+1
View File
@@ -0,0 +1 @@
{"input": {"image": "iVBORw0KGgoAAAANSUhEUgAAAgAAAAIACAIAAAB7GkOtAAAFp0lEQVR42u3VMREAMAgAMfRVRO2woK1GagIJCCB3UfDLR74PwEIhAYABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAIABqABgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgBgACoAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGACAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGACAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAwDiAWweAhQwAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAwAAkADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADADAAFQAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAMQAUAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADADAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADADAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAYNQ/e2eZ8ltazgAAAABJRU5ErkJggg=="}}
+25
View File
@@ -0,0 +1,25 @@
# Combined GPU diffusion worker: FLUX.2 img2img + SDXL inpaint.
# Slim Python base + torch cu124 (torch bundles its own CUDA runtime/cuDNN, so we
# skip the heavy nvidia/cuda base — that halves the image so it fits the worker's
# container disk). The host NVIDIA driver is mounted by RunPod; torch supplies the
# rest. No CUDA-version label = runs on any recent driver. Build ctx = repo root.
FROM python:3.11-slim
ENV DEBIAN_FRONTEND=noninteractive \
PYTHONUNBUFFERED=1 \
PIP_NO_CACHE_DIR=1
RUN apt-get update && apt-get install -y --no-install-recommends \
git libgomp1 libgl1 libglib2.0-0 && \
rm -rf /var/lib/apt/lists/*
WORKDIR /app
RUN pip install --upgrade pip && \
pip install torch --index-url https://download.pytorch.org/whl/cu124
COPY workers/gpu-diffusion/requirements.txt .
RUN pip install -r requirements.txt
COPY workers/gpu-diffusion/handler.py .
CMD ["python3", "-u", "handler.py"]
+148
View File
@@ -0,0 +1,148 @@
"""
Combined GPU diffusion RunPod worker. Routes on input.task:
"img2img" -> FLUX.2 [klein] (prompt / colorize / sky) -> {"image": b64}
"inpaint" -> SDXL inpainting (inpaint/outpaint/eraser/fill) -> {"image": b64}
One request runs exactly one model (isolated). Needs a 24 GB GPU. torch cu124 +
diffusers. Both pipelines load lazily and are cached.
"""
import base64
import io
import os
_V = "/runpod-volume"
if os.path.isdir(_V):
os.environ.setdefault("HF_HOME", os.path.join(_V, "huggingface"))
import torch # noqa: E402
import runpod # noqa: E402
from PIL import Image # noqa: E402
FLUX_MODEL = os.environ.get("FLUX2_MODEL", "black-forest-labs/FLUX.2-klein-4B")
FLUX_STEPS = int(os.environ.get("FLUX2_STEPS", "6"))
INPAINT_MODEL = os.environ.get("SD_INPAINT_MODEL", "diffusers/stable-diffusion-xl-1.0-inpainting-0.1")
MAX_SIZE = int(os.environ.get("MAX_SIZE", "1024"))
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
_HAS24 = DEVICE == "cuda" and torch.cuda.get_device_properties(0).total_memory / 1e9 >= 20
_flux = None
_inpaint = None
def _place(p):
if _HAS24:
return p.to("cuda")
p.enable_model_cpu_offload()
return p
def _flux_pipe():
global _flux
if _flux is None:
from diffusers import Flux2KleinPipeline
p = Flux2KleinPipeline.from_pretrained(FLUX_MODEL, torch_dtype=torch.bfloat16)
p.set_progress_bar_config(disable=True)
_flux = _place(p)
return _flux
def _inpaint_pipe():
global _inpaint
if _inpaint is None:
from diffusers import AutoPipelineForInpainting
try:
p = AutoPipelineForInpainting.from_pretrained(
INPAINT_MODEL, torch_dtype=torch.float16, variant="fp16"
)
except Exception:
p = AutoPipelineForInpainting.from_pretrained(INPAINT_MODEL, torch_dtype=torch.float16)
p.set_progress_bar_config(disable=True)
_inpaint = _place(p)
return _inpaint
def _b64_img(data, mode="RGB"):
if "," in data:
data = data.split(",", 1)[1]
return Image.open(io.BytesIO(base64.b64decode(data))).convert(mode)
def _to_b64(img):
buf = io.BytesIO()
img.save(buf, format="PNG")
return base64.b64encode(buf.getvalue()).decode("utf-8")
def _fit(img, mult=16):
w, h = img.size
s = min(1.0, MAX_SIZE / max(w, h))
nw = max(mult, (int(w * s) // mult) * mult)
nh = max(mult, (int(h * s) // mult) * mult)
return img.resize((nw, nh), Image.LANCZOS)
def _run_flux(inp):
image_b64 = inp.get("image")
prompt = (inp.get("prompt") or "").strip()
if not image_b64:
return {"error": "Missing 'image'."}
if not prompt:
return {"error": "Missing 'prompt'."}
image = _fit(_b64_img(image_b64), 16)
guidance = 1.0
st = inp.get("strength")
if st is not None:
s = max(0.0, min(1.0, float(st)))
guidance = round(1.0 + s * 3.0, 2)
seed = inp.get("seed")
gen = torch.Generator(device="cpu").manual_seed(int(seed)) if seed is not None else None
res = _flux_pipe()(
prompt=prompt, image=image, height=image.height, width=image.width,
guidance_scale=guidance, num_inference_steps=FLUX_STEPS, generator=gen,
).images[0]
return {"image": _to_b64(res)}
def _run_inpaint(inp):
image_b64 = inp.get("image")
mask_b64 = inp.get("mask")
if not image_b64:
return {"error": "Missing 'image'."}
if not mask_b64:
return {"error": "Missing 'mask'."}
prompt = (inp.get("prompt") or "").strip()
negative = (inp.get("negative_prompt") or "").strip() or None
image = _fit(_b64_img(image_b64), 8)
mask = _b64_img(mask_b64, "L").resize(image.size, Image.NEAREST)
steps = int(inp.get("num_inference_steps") or 35)
guidance = float(inp.get("guidance_scale") or 7.0)
raw = inp.get("strength")
strength = 1.0 if not prompt else (max(0.7, min(1.0, float(raw))) if raw is not None else 0.9)
seed = inp.get("seed")
gen = torch.Generator(device=DEVICE).manual_seed(int(seed)) if seed is not None else None
eff = prompt or "clean, seamless, photorealistic background, natural continuation"
res = _inpaint_pipe()(
prompt=eff, negative_prompt=negative, image=image, mask_image=mask,
height=image.height, width=image.width, num_inference_steps=steps,
guidance_scale=guidance, strength=strength, generator=gen,
).images[0]
return {"image": _to_b64(res)}
def handler(event):
inp = (event or {}).get("input") or {}
task = (inp.get("task") or "").strip().lower()
if not task:
task = "inpaint" if inp.get("mask") else "img2img"
try:
if task == "img2img":
return _run_flux(inp)
if task == "inpaint":
return _run_inpaint(inp)
return {"error": f"Unknown task '{task}' for diffusion endpoint."}
except Exception as exc:
return {"error": f"{type(exc).__name__}: {exc}"}
runpod.serverless.start({"handler": handler})
+9
View File
@@ -0,0 +1,9 @@
# FLUX.2 (Flux2KleinPipeline, from git) + SDXL inpaint (AutoPipelineForInpainting).
diffusers @ git+https://github.com/huggingface/diffusers.git
transformers>=4.49.0
accelerate>=1.2.0
safetensors>=0.4.5
sentencepiece>=0.2.0
protobuf>=4.25.0
Pillow==10.4.0
runpod==1.7.9
+30
View File
@@ -0,0 +1,30 @@
# Combined GPU worker: YOLO detect + Real-ESRGAN upscale + Whisper transcribe.
# CUDA 12.1 + torch cu121 (runs on RTX 3090/4090 driver-550 hosts). Build ctx = repo root.
FROM nvidia/cuda:12.1.1-cudnn8-runtime-ubuntu22.04
ENV DEBIAN_FRONTEND=noninteractive \
PYTHONUNBUFFERED=1 \
PIP_NO_CACHE_DIR=1
RUN apt-get update && apt-get install -y --no-install-recommends \
python3 python3-pip libgl1 libglib2.0-0 && \
ln -sf /usr/bin/python3 /usr/bin/python && \
rm -rf /var/lib/apt/lists/*
WORKDIR /app
RUN pip install --upgrade pip && \
pip install torch==2.1.2 torchvision==0.16.2 --index-url https://download.pytorch.org/whl/cu121
COPY workers/gpu-vision/requirements.txt .
RUN pip install -r requirements.txt
# Fix basicsr's removed torchvision import (the #1 Real-ESRGAN crash) — rewrite it.
RUN F=$(python3 -c "import basicsr, os; print(os.path.join(os.path.dirname(basicsr.__file__), 'data', 'degradations.py'))") && \
sed -i 's/from torchvision.transforms.functional_tensor import rgb_to_grayscale/from torchvision.transforms.functional import rgb_to_grayscale/' "$F"
# YOLO construction model + handler.
COPY workers/gpu-vision/best.pt .
ENV YOLO_MODEL=/app/best.pt
COPY workers/gpu-vision/handler.py .
CMD ["python3", "-u", "handler.py"]
Binary file not shown.
+191
View File
@@ -0,0 +1,191 @@
"""
Combined GPU worker: YOLO detect + Real-ESRGAN upscale + Whisper transcribe.
Routes on input.task:
"detect" -> YOLO (best.pt construction model) -> {"detections": [...]}
"upscale" -> Real-ESRGAN (+ optional GFPGAN face) -> {"image": b64}
"transcribe" -> Whisper (faster-whisper) -> {"transcript": ..., ...}
One request runs exactly one model (isolated). torch cu121 + ultralytics +
realesrgan + faster-whisper. Each model loads lazily and is cached.
"""
import base64
import io
import os
import tempfile
import traceback
import urllib.request
_V = "/runpod-volume"
if os.path.isdir(_V):
os.environ.setdefault("HF_HOME", os.path.join(_V, "huggingface"))
import runpod # noqa: E402
from PIL import Image # noqa: E402
_WEIGHTS_DIR = os.path.join(_V, "weights") if os.path.isdir(_V) else "/app/weights"
os.makedirs(_WEIGHTS_DIR, exist_ok=True)
_yolo = None
_ups = None
_face = None
_whisper = None
_RRDB = "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth"
_GFP = "https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.4.pth"
# ---------- YOLO ----------
def _yolo_model():
global _yolo
if _yolo is None:
from ultralytics import YOLO
_yolo = YOLO(os.environ.get("YOLO_MODEL", "/app/best.pt"))
return _yolo
def _detect(inp):
b64 = inp.get("image")
if not b64:
return {"error": "Missing 'image'."}
if "," in b64:
b64 = b64.split(",", 1)[1]
img = Image.open(io.BytesIO(base64.b64decode(b64))).convert("RGB")
conf = float(inp.get("conf") or os.environ.get("YOLO_CONF", "0.15"))
res = _yolo_model().predict(img, conf=conf, verbose=False)
dets = []
for r in res:
names = r.names
for b in (r.boxes or []):
cls = int(b.cls[0])
x1, y1, x2, y2 = [float(v) for v in b.xyxy[0].tolist()]
dets.append({
"name": str(names[cls]),
"confidence": float(b.conf[0]),
"box": {"x1": x1, "y1": y1, "x2": x2, "y2": y2},
})
return {"detections": dets}
# ---------- Real-ESRGAN upscale ----------
def _dl(url, dst):
if not os.path.isfile(dst):
tmp = dst + ".tmp"
urllib.request.urlretrieve(url, tmp)
os.replace(tmp, dst)
return dst
def _upsampler():
global _ups
if _ups is None:
import torch
from basicsr.archs.rrdbnet_arch import RRDBNet
from realesrgan import RealESRGANer
mp = _dl(_RRDB, os.path.join(_WEIGHTS_DIR, "RealESRGAN_x4plus.pth"))
model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=4)
_ups = RealESRGANer(
scale=4, model_path=mp, model=model, tile=0, tile_pad=10, pre_pad=0,
half=torch.cuda.is_available(), gpu_id=None,
)
return _ups
def _face_enh():
global _face
if _face is None:
from gfpgan import GFPGANer
gp = _dl(_GFP, os.path.join(_WEIGHTS_DIR, "GFPGANv1.4.pth"))
_face = GFPGANer(model_path=gp, upscale=4, arch="clean", channel_multiplier=2, bg_upsampler=_upsampler())
return _face
def _upscale(inp):
import numpy as np
b64 = inp.get("image")
if not b64:
return {"error": "Missing 'image'."}
if "," in b64:
b64 = b64.split(",", 1)[1]
scale = int(inp.get("scale", 4))
if scale not in (2, 4):
scale = 4
face = bool(inp.get("face_enhance", False))
pil = Image.open(io.BytesIO(base64.b64decode(b64))).convert("RGB")
rgb = np.array(pil)
bgr = rgb[:, :, ::-1].copy()
if face:
_, _, ob = _face_enh().enhance(bgr, has_aligned=False, only_center_face=False, paste_back=True)
out = Image.fromarray(ob[:, :, ::-1]).resize((rgb.shape[1] * scale, rgb.shape[0] * scale), Image.LANCZOS)
else:
ob, _ = _upsampler().enhance(bgr, outscale=scale)
out = Image.fromarray(ob[:, :, ::-1])
buf = io.BytesIO()
out.save(buf, format="PNG")
return {"image": base64.b64encode(buf.getvalue()).decode("utf-8")}
# ---------- Whisper transcribe ----------
def _whisper_model():
global _whisper
if _whisper is None:
from faster_whisper import WhisperModel
_whisper = WhisperModel(
os.environ.get("WHISPER_MODEL", "medium"),
device=os.environ.get("WHISPER_DEVICE", "cuda"),
compute_type=os.environ.get("WHISPER_COMPUTE", "float16"),
download_root=os.environ.get("HF_HOME"),
)
return _whisper
def _transcribe(inp):
b64 = inp.get("audio")
if not b64:
return {"error": "Missing 'audio'."}
lang = inp.get("language") or None
want_ts = bool(inp.get("timestamps", False))
data = b64.split(",", 1)[1] if "," in b64 else b64
path = None
try:
with tempfile.NamedTemporaryFile(suffix=".audio", delete=False) as f:
f.write(base64.b64decode(data))
path = f.name
segs, info = _whisper_model().transcribe(path, language=lang, vad_filter=True)
parts, out_segs = [], []
for s in segs:
parts.append(s.text)
if want_ts:
out_segs.append({"text": s.text.strip(), "start_sec": round(s.start, 2), "end_sec": round(s.end, 2)})
out = {"transcript": "".join(parts).strip(), "language": info.language}
if want_ts:
out["segments"] = out_segs
return out
finally:
if path and os.path.exists(path):
os.unlink(path)
def handler(event):
inp = (event or {}).get("input") or {}
task = (inp.get("task") or "").strip().lower()
if not task:
task = "transcribe" if inp.get("audio") else ("upscale" if inp.get("scale") else "detect")
try:
if task == "detect":
return _detect(inp)
if task == "upscale":
return _upscale(inp)
if task == "transcribe":
return _transcribe(inp)
return {"error": f"Unknown task '{task}' for vision endpoint."}
except Exception as exc:
traceback.print_exc()
return {"error": f"{type(exc).__name__}: {exc}"}
runpod.serverless.start({"handler": handler})
+10
View File
@@ -0,0 +1,10 @@
# YOLO (ultralytics) + Real-ESRGAN (realesrgan/basicsr/gfpgan) + Whisper (faster-whisper).
# torch/torchvision are installed in the Dockerfile first (cu121).
runpod==1.7.9
ultralytics==8.3.40
realesrgan==0.3.0
basicsr==1.4.2
gfpgan==1.3.8
faster-whisper==1.0.3
numpy==1.24.4
Pillow==10.4.0
+26
View File
@@ -0,0 +1,26 @@
# RunPod Serverless worker: SDXL inpainting (Magic Eraser / Generative Fill / Outpaint).
# CUDA 12.1 base + torch cu121 — the same base the detect/upscale workers use, which
# runs on RunPod's RTX 3090/4090 (driver 550) hosts (12.1 <= 12.4, backward-compatible).
# Build context = repo root, so COPY paths are workers/inpaint/*.
FROM nvidia/cuda:12.1.1-cudnn8-runtime-ubuntu22.04
ENV DEBIAN_FRONTEND=noninteractive \
PYTHONUNBUFFERED=1 \
PIP_NO_CACHE_DIR=1
RUN apt-get update && apt-get install -y --no-install-recommends \
python3 python3-pip git && \
ln -sf /usr/bin/python3 /usr/bin/python && \
rm -rf /var/lib/apt/lists/*
WORKDIR /app
# torch (CUDA 12.1 build) FIRST — carries sm_86 (3090) + sm_89 (4090) kernels.
RUN pip install --upgrade pip && \
pip install torch --index-url https://download.pytorch.org/whl/cu121
COPY workers/inpaint/requirements.txt .
RUN pip install -r requirements.txt
COPY workers/inpaint/handler.py .
CMD ["python3", "-u", "handler.py"]
+30
View File
@@ -0,0 +1,30 @@
# Inpaint worker — SDXL (`diffusers/stable-diffusion-xl-1.0-inpainting-0.1`)
One RunPod serverless endpoint that powers **Magic Eraser**, **Generative Fill**,
and **Outpaint** in the photo editor. Set its `/runsync` URL as the Vercel env var
**`RUNPOD_SD_INPAINT_URL`**.
## Contract
Request `{"input": {...}}` — matches the app's `rpInpaint`:
| field | required | meaning |
|---|---|---|
| `image` | yes | base64 image (raw or `data:` URI) |
| `mask` | yes | base64 mask — **white = regenerate, black = keep** |
| `prompt` | no | `""` = Magic Eraser (clean fill); text = Generative Fill / Outpaint |
| `strength` | no | 0..1 (empty prompt forces 1.0 to fully erase) |
| `guidance_scale` | no | default 7 |
| `num_inference_steps` | no | default 35 |
| `negative_prompt` | no | — |
| `seed` | no | — |
Response: `{"image": "<base64 PNG>"}` or `{"error": "..."}`.
## RunPod setup
1. New Serverless endpoint → connect this GitHub repo → **Dockerfile path** `workers/inpaint/Dockerfile` (build context = repo root).
2. GPU: **24 GB** (RTX 4090 / 3090). Max workers ≥ 1; optionally a network volume to cache the ~7 GB model so cold starts don't re-download.
3. Copy the endpoint's `…/runsync` URL → Vercel `RUNPOD_SD_INPAINT_URL`, then redeploy the web app.
## Env knobs
- `SD_INPAINT_MODEL` (default `diffusers/stable-diffusion-xl-1.0-inpainting-0.1`)
- `SD_INPAINT_MAX_SIZE` (default `1024`)
+137
View File
@@ -0,0 +1,137 @@
"""
RunPod Serverless worker: SDXL inpainting.
Powers the app's Magic Eraser, Generative Fill AND Outpaint — all three go
through this one endpoint (RUNPOD_SD_INPAINT_URL). No app change needed; this
matches the exact contract the app already sends (rpInpaint):
{"input": {
"image": "<base64>", # required (raw base64 or data: URI)
"mask": "<base64>", # required — WHITE = regenerate, BLACK = keep
"prompt": "<text>", # optional ("" = Magic Eraser: clean fill)
"strength": 0.8, # optional
"guidance_scale": 7, # optional
"num_inference_steps": 35, # optional
"negative_prompt": "...", # optional
"seed": 123 # optional
}}
Returns: {"image": "<base64 PNG>"} (or {"error": "..."}).
CUDA 12.1 base + torch cu121 (same as the detect/upscale workers) so it runs on
RunPod's RTX 3090/4090 (driver 550) hosts.
"""
import base64
import io
import os
_V = "/runpod-volume"
if os.path.isdir(_V):
os.environ.setdefault("HF_HOME", os.path.join(_V, "huggingface"))
import torch # noqa: E402
import runpod # noqa: E402
from PIL import Image # noqa: E402
from diffusers import AutoPipelineForInpainting # noqa: E402
MODEL = os.environ.get("SD_INPAINT_MODEL", "diffusers/stable-diffusion-xl-1.0-inpainting-0.1")
MAX_SIZE = int(os.environ.get("SD_INPAINT_MAX_SIZE", "1024"))
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
_pipe = None
def _get_pipe():
global _pipe
if _pipe is None:
try:
pipe = AutoPipelineForInpainting.from_pretrained(
MODEL, torch_dtype=torch.float16, variant="fp16"
)
except Exception:
# Repo may not ship an fp16 variant — fall back to default weights.
pipe = AutoPipelineForInpainting.from_pretrained(MODEL, torch_dtype=torch.float16)
total_gb = torch.cuda.get_device_properties(0).total_memory / 1e9 if DEVICE == "cuda" else 0
if total_gb >= 20:
pipe = pipe.to("cuda")
else:
pipe.enable_model_cpu_offload()
pipe.set_progress_bar_config(disable=True)
_pipe = pipe
return _pipe
def _b64_to_image(data, mode="RGB"):
if "," in data:
data = data.split(",", 1)[1]
return Image.open(io.BytesIO(base64.b64decode(data))).convert(mode)
def _image_to_b64(img):
buf = io.BytesIO()
img.save(buf, format="PNG")
return base64.b64encode(buf.getvalue()).decode("utf-8")
def _fit(img):
# SDXL wants multiples of 8; cap the long side at MAX_SIZE.
w, h = img.size
scale = min(1.0, MAX_SIZE / max(w, h))
nw = max(8, (int(w * scale) // 8) * 8)
nh = max(8, (int(h * scale) // 8) * 8)
return img.resize((nw, nh), Image.LANCZOS)
def handler(event):
inp = (event or {}).get("input") or {}
image_b64 = inp.get("image")
mask_b64 = inp.get("mask")
if not image_b64:
return {"error": "Missing 'image' (base64)."}
if not mask_b64:
return {"error": "Missing 'mask' (base64)."}
prompt = (inp.get("prompt") or "").strip()
negative = (inp.get("negative_prompt") or "").strip() or None
try:
image = _fit(_b64_to_image(image_b64, "RGB"))
mask = _b64_to_image(mask_b64, "L").resize(image.size, Image.NEAREST)
steps = int(inp.get("num_inference_steps") or 35)
guidance = float(inp.get("guidance_scale") or 7.0)
# Strength: empty prompt = Magic Eraser -> fully replace (1.0). With a
# prompt (Generative Fill / Outpaint) respect the caller but keep a floor
# so the masked region actually changes.
raw = inp.get("strength")
if not prompt:
strength = 1.0
else:
strength = max(0.7, min(1.0, float(raw))) if raw is not None else 0.9
seed = inp.get("seed")
generator = (
torch.Generator(device=DEVICE).manual_seed(int(seed)) if seed is not None else None
)
eff_prompt = prompt or "clean, seamless, photorealistic background, natural continuation"
result = _get_pipe()(
prompt=eff_prompt,
negative_prompt=negative,
image=image,
mask_image=mask,
height=image.height,
width=image.width,
num_inference_steps=steps,
guidance_scale=guidance,
strength=strength,
generator=generator,
).images[0]
return {"image": _image_to_b64(result)}
except Exception as exc:
return {"error": f"{type(exc).__name__}: {exc}"}
runpod.serverless.start({"handler": handler})
+8
View File
@@ -0,0 +1,8 @@
# SDXL inpainting worker. torch is installed (cu121) in the Dockerfile first.
# AutoPipelineForInpainting + SDXL uses two text encoders → transformers.
diffusers>=0.30.0
transformers>=4.44.0
accelerate>=0.34.0
safetensors>=0.4.5
Pillow==10.4.0
runpod==1.7.9
+1
View File
@@ -0,0 +1 @@
{"input": {"image": "iVBORw0KGgoAAAANSUhEUgAAAgAAAAIACAIAAAB7GkOtAAAFlklEQVR42u3VMQEAAAzCMKQjHQ97l0jo0xSAlyIBgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGACAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGACAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGACAAUgAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGACAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGACAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAADcDrctaAb6XeXAAAAAASUVORK5CYII=", "mask": "iVBORw0KGgoAAAANSUhEUgAAAgAAAAIACAAAAADRE4smAAADXUlEQVR42u3SQREAAAzCMPybhvc0LJXQS6rXxQIABIAAEAACQAAIAAEgAASAABAAAkAACAABIAAEgAAQAAJAAAgAASAABIAAEAACQAAIAAEgAASAABAAAkAACAABIAAEgAAQAAJAAAgAASAABIAAEAACQAAIAAEgAASAABAAAkAACAABIAAEgAAQAAJAAAgAASAABIAAEAACQAAIAAEgAASAABAAAgAACwAQAAJAAAgAASAABIAAEAACQAAIAAEgAASAABAAAkAACAABIAAEgAAQAAJAAAgAASAABIAAEAACQAAIAAEgAASAABAAAkAACAABIAAEgAAQAAJAAAgAASAABIAAEAACQAAIAAEgAASAABAAAkAACAABIAAEgAAQAAJAAAgAASAABIAAEAACQAAIAAAsAEAACAABIAAEgAAQAAJAAAgAASAABIAAEAACQAAIAAEgAASAABAAAkAACAABIAAEgAAQAAJAAAgAASAABIAAEAACQAAIAAEgAASAABAAAkAACAABIAAEgAAQAAJAAAgAASAABIAAEAACQAAIAAEgAASAABAAAkAACAABIAAEgAAQAAJAAAgAASAAALAAAAEgAASAABAAAkAACAABIAAEgAAQAAJAAAgAASAABIAAEAACQAAIAAEgAASAABAAAkAACAABIAAEgAAQAAJAAAgAASAABIAAEAACQAAIAAEgAASAABAAAkAACAABIAAEgAAQAAJAAAgAASAABIAAEAACQAAIAAEgAASAABAAAkAACAABIAAEgAAAwAIABIAAEAACQAAIAAEgAASAABAAAkAACAABIAAEgAAQAAJAAAgAASAABIAAEAACQAAIAAEgAASAABAAAkAACAABIAAEgAAQAAJAAAgAASAABIAAEAACQAAIAAEgAASAABAAAkAACAABIAAEgAAQAAJAAAgAASAABIAAEAACQAAIAAEgAASAABAAAgAACwAQAAJAAAgAASAABIAAEAACQAAIAAEgAASAABAAAkAACAABIAAEgAAQAAJAAAgAASAABIAAEAACQAAIAAEgAASAABAAAkAACAABIAAEgAAQAAJAAAgAASAABIAAEAACQAAIAAEgAASAABAAAkAACAABIAAEgAAQAAJAAAgAASAABIBuA8sSO8WbTgaBAAAAAElFTkSuQmCC", "prompt": "a clean concrete wall", "num_inference_steps": 20}}
+33
View File
@@ -0,0 +1,33 @@
FROM nvidia/cuda:12.1.1-cudnn8-runtime-ubuntu22.04
ENV DEBIAN_FRONTEND=noninteractive \
PYTHONUNBUFFERED=1 \
PIP_NO_CACHE_DIR=1
RUN apt-get update && apt-get install -y --no-install-recommends \
python3 python3-pip git ca-certificates \
libgl1 libglib2.0-0 && \
rm -rf /var/lib/apt/lists/* && \
ln -sf /usr/bin/python3 /usr/bin/python
WORKDIR /app
# Clone the research repo (inference code + configs). Pinning is recommended:
# replace HEAD with a specific commit once verified.
RUN git clone --depth 1 https://github.com/AlanSavio25/DeepSingleImageCalibration.git \
/opt/DeepSingleImageCalibration
# Install the CUDA 12.1 torch stack + repo deps.
COPY workers/tilt/requirements.txt .
RUN pip install --index-url https://download.pytorch.org/whl/cu121 \
torch==2.1.2 torchvision==0.16.2 && \
pip install -r requirements.txt
# Best-effort install of the repo's own requirements if present (non-fatal).
RUN if [ -f /opt/DeepSingleImageCalibration/requirements.txt ]; then \
pip install -r /opt/DeepSingleImageCalibration/requirements.txt || true ; \
fi
COPY workers/tilt/handler.py .
CMD ["python3", "-u", "handler.py"]
+35
View File
@@ -0,0 +1,35 @@
# workers/tilt — Single-Image Camera Calibration (roll / pitch / FOV)
Predicts camera **roll**, **pitch**, and **vertical field-of-view** from one image, based on
[AlanSavio25/DeepSingleImageCalibration](https://github.com/AlanSavio25/DeepSingleImageCalibration).
- **Dockerfile:** `/workers/tilt/Dockerfile` (build context = repo root)
- **GPU:** light — NVIDIA **T4** (16 GB) is plenty; CPU also works, slower.
- **App env var to set on the endpoint:** `RUNPOD_TILT_URL` — the deployed RunPod endpoint URL the app calls.
## Contract
Input:
```json
{ "input": { "image": "<base64>" } }
```
`image` may include a `data:image/...;base64,` prefix (stripped automatically).
Output:
```json
{ "roll_degrees": 0.0, "pitch_degrees": 0.0, "fov_degrees": 0.0 }
```
On error: `{ "error": "..." }`.
## Weights (required — research repo, not on pip)
The pretrained checkpoint is **not** bundled. Download the repo's released weights and place
them on the network volume, then point the worker at them:
- Default path: `/runpod-volume/tilt/weights.tar`
- Override with env var `WEIGHTS_PATH`.
The backbone loads into the volume-cached `TORCH_HOME`/`HF_HOME` on first request.
## Tunable env vars (no rebuild needed)
`WEIGHTS_PATH`, `CALIB_BACKBONE` (densenet161|densenet121), `CALIB_NUM_BINS`,
`CALIB_INPUT_SIZE`, `CALIB_ROLL_MIN/MAX`, `CALIB_VFOV_MIN/MAX`, `CALIB_RHO_MIN/MAX`,
`CALIB_DIRECT_PITCH` (=1 if the net outputs pitch directly), `CALIB_PITCH_MIN/MAX`.
+206
View File
@@ -0,0 +1,206 @@
# --- network volume cache setup (must run before importing heavy libs) ---
_V = "/runpod-volume"
import os
if os.path.isdir(_V):
os.environ.setdefault("HF_HOME", os.path.join(_V, "huggingface"))
os.environ.setdefault("TORCH_HOME", os.path.join(_V, "torch"))
import io
import sys
import base64
import math
import runpod
# The DeepSingleImageCalibration repo is cloned into /opt/DeepSingleImageCalibration
# by the Dockerfile. Make it importable.
_REPO = os.environ.get("CALIB_REPO", "/opt/DeepSingleImageCalibration")
if _REPO not in sys.path:
sys.path.insert(0, _REPO)
_MODEL = None # cached (model, device)
_TF = None # cached torchvision transform
# ---------------------------------------------------------------------------
# Bin / range configuration.
#
# NOTE (RESEARCH REPO -- VERIFY THESE): AlanSavio25/DeepSingleImageCalibration
# predicts each parameter as a classification over discretised bins, then takes
# the soft-expected value over the bin centres. The exact number of bins, the
# min/max of each range, the head ORDER/NAMES, and whether the network emits
# `pitch` directly or a horizon-offset `rho` all come from the repo's training
# config + the single-image-prediction notebook. The values below follow the
# "A Perceptual Measure for Deep Single Image Camera Calibration" formulation
# and are the most likely defaults, but MUST be confirmed against the checkpoint.
# Override any of them via env vars without rebuilding.
# ---------------------------------------------------------------------------
NUM_BINS = int(os.environ.get("CALIB_NUM_BINS", "256"))
ROLL_MIN = float(os.environ.get("CALIB_ROLL_MIN", "-45.0"))
ROLL_MAX = float(os.environ.get("CALIB_ROLL_MAX", "45.0"))
VFOV_MIN = float(os.environ.get("CALIB_VFOV_MIN", "20.0"))
VFOV_MAX = float(os.environ.get("CALIB_VFOV_MAX", "105.0"))
# rho = signed vertical horizon offset as a fraction of image height.
RHO_MIN = float(os.environ.get("CALIB_RHO_MIN", "-1.5"))
RHO_MAX = float(os.environ.get("CALIB_RHO_MAX", "1.5"))
# If the network already emits pitch (degrees) instead of rho, set this to "1".
CALIB_DIRECT_PITCH = os.environ.get("CALIB_DIRECT_PITCH", "0") == "1"
PITCH_MIN = float(os.environ.get("CALIB_PITCH_MIN", "-45.0"))
PITCH_MAX = float(os.environ.get("CALIB_PITCH_MAX", "45.0"))
INPUT_SIZE = int(os.environ.get("CALIB_INPUT_SIZE", "320"))
BACKBONE = os.environ.get("CALIB_BACKBONE", "densenet161") # repo default backbone
# Where the pretrained checkpoint lives. Download the repo's released weights
# onto the network volume and point this at them.
WEIGHTS_PATH = os.environ.get(
"WEIGHTS_PATH",
os.path.join(_V, "tilt", "weights.tar") if os.path.isdir(_V) else "/opt/weights.tar",
)
def _build_model():
import torch
import torch.nn as nn
import torchvision
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# DenseNet backbone with one classification head per predicted parameter.
# This mirrors the repo architecture (ImageNet-pretrained backbone whose
# classifier is replaced by per-parameter heads over NUM_BINS bins).
if BACKBONE == "densenet161":
base = torchvision.models.densenet161(weights=None)
feat_dim = 2208
elif BACKBONE == "densenet121":
base = torchvision.models.densenet121(weights=None)
feat_dim = 1024
else:
raise RuntimeError(f"Unsupported CALIB_BACKBONE={BACKBONE!r}")
class CalibNet(nn.Module):
def __init__(self):
super().__init__()
self.features = base.features
self.head_roll = nn.Linear(feat_dim, NUM_BINS)
self.head_pitch = nn.Linear(feat_dim, NUM_BINS) # 'rho' or pitch head
self.head_vfov = nn.Linear(feat_dim, NUM_BINS)
self.head_k1 = nn.Linear(feat_dim, NUM_BINS) # distortion (unused in output)
def forward(self, x):
f = self.features(x)
f = nn.functional.relu(f, inplace=True)
f = nn.functional.adaptive_avg_pool2d(f, (1, 1)).flatten(1)
return {
"roll": self.head_roll(f),
"pitch": self.head_pitch(f),
"vfov": self.head_vfov(f),
"k1": self.head_k1(f),
}
model = CalibNet()
# Load pretrained weights. Research checkpoints commonly store the state dict
# under a "model" or "state_dict" key; fall back to the raw object.
if not os.path.isfile(WEIGHTS_PATH):
raise RuntimeError(
f"checkpoint not found at {WEIGHTS_PATH!r}; place the repo's pretrained "
f"weights on the network volume or set WEIGHTS_PATH"
)
ckpt = torch.load(WEIGHTS_PATH, map_location="cpu")
if isinstance(ckpt, dict):
state = ckpt.get("model", ckpt.get("state_dict", ckpt))
else:
state = ckpt
if isinstance(state, dict):
# strip common prefixes
cleaned = {}
for k, v in state.items():
nk = k
for pref in ("module.", "model.", "net."):
if nk.startswith(pref):
nk = nk[len(pref):]
cleaned[nk] = v
# strict=False: head naming in the checkpoint may differ from ours; the
# backbone (bulk of the weights) will still load. See notes.
model.load_state_dict(cleaned, strict=False)
model.eval().to(device)
return model, device
def _get_transform():
global _TF
if _TF is None:
import torchvision.transforms as T
_TF = T.Compose([
T.Resize((INPUT_SIZE, INPUT_SIZE)),
T.ToTensor(),
T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])
return _TF
def _load_model():
global _MODEL
if _MODEL is None:
_MODEL = _build_model()
return _MODEL
def _expected_value(logits, vmin, vmax):
"""Soft-argmax: softmax over bins, dot with evenly-spaced bin centres."""
import torch
probs = torch.softmax(logits, dim=-1)
n = probs.shape[-1]
centers = torch.linspace(vmin, vmax, n, device=probs.device, dtype=probs.dtype)
return float((probs * centers).sum(dim=-1).item())
def _decode_image(b64):
from PIL import Image
if "," in b64:
b64 = b64.split(",", 1)[1]
raw = base64.b64decode(b64)
return Image.open(io.BytesIO(raw)).convert("RGB")
def handler(event):
try:
data = (event or {}).get("input") or {}
b64 = data.get("image")
if not b64:
return {"error": "missing required 'image' (base64) field"}
import torch
img = _decode_image(b64)
model, device = _load_model()
tf = _get_transform()
x = tf(img).unsqueeze(0).to(device)
with torch.no_grad():
out = model(x)
roll = _expected_value(out["roll"], ROLL_MIN, ROLL_MAX)
vfov = _expected_value(out["vfov"], VFOV_MIN, VFOV_MAX)
if CALIB_DIRECT_PITCH:
pitch = _expected_value(out["pitch"], PITCH_MIN, PITCH_MAX)
else:
# Head predicts rho = signed vertical horizon offset as a fraction of
# image height. Convert to pitch via the pinhole relation:
# tan(pitch) = 2 * rho * tan(vfov/2)
rho = _expected_value(out["pitch"], RHO_MIN, RHO_MAX)
vfov_rad = math.radians(vfov)
pitch = math.degrees(math.atan(2.0 * rho * math.tan(vfov_rad / 2.0)))
return {
"roll_degrees": float(roll),
"pitch_degrees": float(pitch),
"fov_degrees": float(vfov),
}
except Exception as exc:
return {"error": f"{type(exc).__name__}: {exc}"}
runpod.serverless.start({"handler": handler})
+9
View File
@@ -0,0 +1,9 @@
runpod==1.7.9
numpy==1.26.4
Pillow==10.4.0
opencv-python-headless==4.10.0.84
omegaconf==2.3.0
tqdm==4.66.5
matplotlib==3.8.4
h5py==3.11.0
scipy==1.13.1
+1
View File
@@ -0,0 +1 @@
{"input": {"image": "iVBORw0KGgoAAAANSUhEUgAAAgAAAAIACAIAAAB7GkOtAAAFp0lEQVR42u3VMREAMAgAMfRVRO2woK1GagIJCCB3UfDLR74PwEIhAYABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAIABqABgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgBgACoAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGACAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGACAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAwDiAWweAhQwAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAwAAkADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADADAAFQAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAMQAUAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADADAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADADAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAYNQ/e2eZ8ltazgAAAABJRU5ErkJggg=="}}
+34
View File
@@ -0,0 +1,34 @@
FROM nvidia/cuda:12.1.1-cudnn8-runtime-ubuntu22.04
ENV DEBIAN_FRONTEND=noninteractive
ENV PIP_NO_CACHE_DIR=1
ENV PYTHONUNBUFFERED=1
RUN apt-get update && apt-get install -y --no-install-recommends \
python3 python3-pip python3-dev \
libgl1 libglib2.0-0 wget ca-certificates \
&& rm -rf /var/lib/apt/lists/* \
&& ln -sf /usr/bin/python3 /usr/bin/python
WORKDIR /app
# Newer CUDA-12.1 torch so the compiled kernels cover current GPUs (fixes
# "no kernel image is available for execution on the device"). The basicsr
# functional_tensor import (removed in torchvision>=0.16) is handled by the
# sed-patch below instead of pinning an old torchvision.
RUN pip install --upgrade pip && \
pip install torch==2.1.2 torchvision==0.16.2 \
--index-url https://download.pytorch.org/whl/cu121
COPY workers/upscale/requirements.txt .
RUN pip install -r requirements.txt
# Robustly fix basicsr's removed torchvision import (the #1 Real-ESRGAN crash).
# Works whether or not the torch/torchvision pin held — rewrites the bad import.
RUN F=$(python3 -c 'import basicsr,os;print(os.path.join(os.path.dirname(basicsr.__file__),"data","degradations.py"))') && \
sed -i 's/from torchvision.transforms.functional_tensor import rgb_to_grayscale/from torchvision.transforms.functional import rgb_to_grayscale/' "$F" && \
echo "patched basicsr degradations.py"
COPY workers/upscale/handler.py .
CMD ["python3","-u","handler.py"]
+35
View File
@@ -0,0 +1,35 @@
# Real-ESRGAN Upscale Worker (#7)
RunPod Serverless worker for Real-ESRGAN image super-resolution with optional GFPGAN face enhancement.
- **Dockerfile path:** `/workers/upscale/Dockerfile`
- **Build context:** repository ROOT (COPY paths are prefixed with `workers/upscale/`).
- **GPU:** medium (RTX 3090 / RTX 4090, ~24 GB VRAM).
- **App env var to set:** `RUNPOD_UPSCALE_URL` (point your app at this endpoint's RunPod URL).
## Contract
Input:
```json
{"image": "<base64 PNG/JPEG, optional data: prefix>", "scale": 2 or 4 (default 4), "face_enhance": false}
```
Output (success):
```json
{"image": "<base64 PNG>"}
```
Output (error):
```json
{"error": "SomeError: message"}
```
## Model weights
Downloaded on first request and cached to the network volume when mounted:
- `RealESRGAN_x4plus.pth` (RRDBNet x4) -> `/runpod-volume/weights/`
- `GFPGANv1.4.pth` (only when `face_enhance=true`) -> `/runpod-volume/weights/`
Attach a network volume mounted at `/runpod-volume` so weights (and `HF_HOME`) persist between cold starts. Without a volume, weights download to `/app/weights` per container.
## Notes
The base model is x4; `scale=2` is produced via RealESRGANer's `outscale` parameter. For `face_enhance`, GFPGAN upscales at its fixed factor and the result is resized to the requested scale.
+113
View File
@@ -0,0 +1,113 @@
_V = "/runpod-volume"
import os
if os.path.isdir(_V):
os.environ.setdefault("HF_HOME", os.path.join(_V, "huggingface"))
import io
import base64
import traceback
import urllib.request
import numpy as np
from PIL import Image
_WEIGHTS_DIR = os.path.join(_V, "weights") if os.path.isdir(_V) else "/app/weights"
os.makedirs(_WEIGHTS_DIR, exist_ok=True)
_RRDB_URL = "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth"
_GFPGAN_URL = "https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.4.pth"
_UPSAMPLER = None
_FACE_ENHANCER = None
def _download(url, dst):
if not os.path.isfile(dst):
tmp = dst + ".tmp"
urllib.request.urlretrieve(url, tmp)
os.replace(tmp, dst)
return dst
def _get_upsampler():
global _UPSAMPLER
if _UPSAMPLER is None:
import torch
from basicsr.archs.rrdbnet_arch import RRDBNet
from realesrgan import RealESRGANer
model_path = _download(_RRDB_URL, os.path.join(_WEIGHTS_DIR, "RealESRGAN_x4plus.pth"))
model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=4)
_UPSAMPLER = RealESRGANer(
scale=4,
model_path=model_path,
model=model,
tile=0,
tile_pad=10,
pre_pad=0,
half=torch.cuda.is_available(),
gpu_id=None,
)
return _UPSAMPLER
def _get_face_enhancer():
global _FACE_ENHANCER
if _FACE_ENHANCER is None:
from gfpgan import GFPGANer
gfpgan_path = _download(_GFPGAN_URL, os.path.join(_WEIGHTS_DIR, "GFPGANv1.4.pth"))
_FACE_ENHANCER = GFPGANer(
model_path=gfpgan_path,
upscale=4,
arch="clean",
channel_multiplier=2,
bg_upsampler=_get_upsampler(),
)
return _FACE_ENHANCER
def handler(event):
try:
data = (event or {}).get("input") or {}
b64 = data.get("image")
if not b64:
return {"error": "ValueError: 'image' is required"}
if "," in b64:
b64 = b64.split(",", 1)[1]
scale = int(data.get("scale", 4))
if scale not in (2, 4):
scale = 4
face_enhance = bool(data.get("face_enhance", False))
raw = base64.b64decode(b64)
pil = Image.open(io.BytesIO(raw)).convert("RGB")
rgb = np.array(pil)
bgr = rgb[:, :, ::-1].copy() # RealESRGANer expects BGR (cv2 convention)
if face_enhance:
enhancer = _get_face_enhancer()
_, _, out_bgr = enhancer.enhance(
bgr, has_aligned=False, only_center_face=False, paste_back=True
)
# GFPGANer upscales by its fixed factor; resize to the requested scale.
target = (rgb.shape[1] * scale, rgb.shape[0] * scale)
out_rgb = out_bgr[:, :, ::-1]
out_pil = Image.fromarray(out_rgb).resize(target, Image.LANCZOS)
else:
upsampler = _get_upsampler()
out_bgr, _ = upsampler.enhance(bgr, outscale=scale)
out_rgb = out_bgr[:, :, ::-1]
out_pil = Image.fromarray(out_rgb)
buf = io.BytesIO()
out_pil.save(buf, format="PNG")
return {"image": base64.b64encode(buf.getvalue()).decode("utf-8")}
except Exception as exc:
traceback.print_exc()
return {"error": f"{type(exc).__name__}: {exc}"}
import runpod
runpod.serverless.start({"handler": handler})
+6
View File
@@ -0,0 +1,6 @@
runpod==1.7.9
realesrgan==0.3.0
basicsr==1.4.2
gfpgan==1.3.8
numpy==1.24.4
pillow==10.4.0
+1
View File
@@ -0,0 +1 @@
{"input": {"image": "iVBORw0KGgoAAAANSUhEUgAAAgAAAAIACAIAAAB7GkOtAAAFp0lEQVR42u3VMREAMAgAMfRVRO2woK1GagIJCCB3UfDLR74PwEIhAYABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAIABqABgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgBgACoAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGACAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGACAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAwDiAWweAhQwAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAwAAkADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADADAAFQAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAMQAUAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADADAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADADAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAYNQ/e2eZ8ltazgAAAABJRU5ErkJggg==", "scale": 4}}
+20
View File
@@ -0,0 +1,20 @@
# RunPod Serverless worker: Speech-to-Text (faster-whisper).
# Lightweight — no torch/NeMo; ctranslate2 uses the CUDA + cuDNN in the base image.
# Paths are relative to the REPO ROOT (RunPod's build context).
FROM nvidia/cuda:12.1.1-cudnn8-runtime-ubuntu22.04
ENV DEBIAN_FRONTEND=noninteractive \
PYTHONUNBUFFERED=1 \
PIP_NO_CACHE_DIR=1
RUN apt-get update && apt-get install -y --no-install-recommends \
python3 python3-pip ffmpeg && \
ln -sf /usr/bin/python3 /usr/bin/python && \
rm -rf /var/lib/apt/lists/*
WORKDIR /app
COPY workers/voice-to-text/requirements.txt .
RUN pip3 install --no-cache-dir -r requirements.txt
COPY workers/voice-to-text/handler.py .
CMD ["python3", "-u", "handler.py"]
+48
View File
@@ -0,0 +1,48 @@
# Speech-to-Text worker — Whisper (RunPod Serverless)
Multilingual voice-to-text via **faster-whisper** — lets a worker **speak**
instead of typing annotation text, in both the image and video editors. 99
languages incl. Hindi/regional + English. Chosen over Parakeet/NeMo because it
has no heavy dependencies (no NeMo, no torch), so it builds cleanly.
## Contract
Request → `POST /runsync`:
```json
{ "input": { "audio": "<base64 audio>", "language": "en", "timestamps": false } }
```
- `audio`: base64 wav/mp3/m4a (any sample rate — resampled internally)
- `language`: optional (e.g. `"hi"`, `"en"`); omit to auto-detect
- `timestamps`: optional — per-segment start/end
Response → `{ "output": { "transcript": "...", "language": "en" } }`
(with `segments: [{text, start_sec, end_sec}]` if `timestamps: true`), or `{ "output": { "error": "..." } }`.
## Deploy (new endpoint — one per model)
1. **RunPod → Serverless → New Endpoint → Deploy from a GitHub repository**`Sumit-Pluto/photo_gallery`.
2. **Dockerfile path:** `/workers/voice-to-text/Dockerfile`
3. Config:
- **GPU:** light — **16 GB (T4 / A4000)** is plenty (`large-v3` ≈ 5 GB VRAM).
- **Network volume:** attach your weights volume (caches the model).
- **Workers:** Max 1, Active 0, FlashBoot on. **Execution timeout** ~300s (first-run download).
4. **Warm up once** (downloads the model, ~1.5 GB):
```bash
curl -X POST https://api.runpod.ai/v2/<NEW_ID>/runsync \
-H "Content-Type: application/json" -H "Authorization: Bearer <API_KEY>" \
-d '{"input":{"audio":"<base64-wav>","language":"en"}}'
```
5. In **Vercel**, set `RUNPOD_STT_URL` = `https://api.runpod.ai/v2/<NEW_ID>/runsync` (server-only).
## Tuning (env vars)
| Var | Default | Purpose |
|---|---|---|
| `WHISPER_MODEL` | `large-v3` | accuracy vs speed — `medium` / `small` are faster/lighter |
| `WHISPER_COMPUTE` | `float16` | `int8_float16` uses less VRAM |
## App integration (next step, not yet built)
Mic-record button in the editor's text/annotation tool → `/api/ai/transcribe`
(proxies here) → inserts the transcript. Client `rpTranscribe()` already exists in
`apps/web/src/lib/runpod/endpoints.ts`.
+91
View File
@@ -0,0 +1,91 @@
"""
RunPod Serverless worker: Speech-to-Text (Whisper via faster-whisper).
Chosen over Parakeet/NeMo because it has NO heavy dependencies (no NeMo, no
torch — just ctranslate2), so it builds cleanly, AND it's multilingual (99
languages incl. Hindi/regional + English) — a better fit for a mixed crew.
Lets a worker speak instead of typing annotation text (image + video editors).
Matches the app's transcribe contract:
{"input": {
"audio": "<base64 audio>", # required (wav/mp3/m4a; any sample rate — resampled internally)
"language": "en", # optional (e.g. "hi", "en"); omit = auto-detect
"timestamps": false # optional — return per-segment start/end
}}
Returns: {"transcript": "...", "language": "en", "segments"?: [{text,start_sec,end_sec}]}
(or {"error": "..."}).
Tuning (env): WHISPER_MODEL (default large-v3; "medium"/"small" are faster),
WHISPER_COMPUTE (default float16; "int8_float16" for less VRAM).
"""
import base64
import os
import tempfile
_VOLUME = "/runpod-volume"
if os.path.isdir(_VOLUME):
os.environ.setdefault("HF_HOME", os.path.join(_VOLUME, "huggingface"))
import runpod # noqa: E402
from faster_whisper import WhisperModel # noqa: E402
MODEL_SIZE = os.environ.get("WHISPER_MODEL", "large-v3")
DEVICE = os.environ.get("WHISPER_DEVICE", "cuda")
COMPUTE = os.environ.get("WHISPER_COMPUTE", "float16")
_model = None
def _get_model():
global _model
if _model is None:
_model = WhisperModel(
MODEL_SIZE,
device=DEVICE,
compute_type=COMPUTE,
download_root=os.environ.get("HF_HOME"),
)
return _model
def handler(event):
inp = (event or {}).get("input") or {}
audio_b64 = inp.get("audio")
if not audio_b64:
return {"error": "Missing 'audio' (base64)."}
language = inp.get("language") or None
want_ts = bool(inp.get("timestamps", False))
path = None
try:
data = audio_b64.split(",", 1)[1] if "," in audio_b64 else audio_b64
with tempfile.NamedTemporaryFile(suffix=".audio", delete=False) as f:
f.write(base64.b64decode(data))
path = f.name
segments, info = _get_model().transcribe(path, language=language, vad_filter=True)
parts = []
segs = []
for s in segments:
parts.append(s.text)
if want_ts:
segs.append(
{"text": s.text.strip(), "start_sec": round(s.start, 2), "end_sec": round(s.end, 2)}
)
out = {"transcript": "".join(parts).strip(), "language": info.language}
if want_ts:
out["segments"] = segs
return out
except Exception as exc:
return {"error": f"{type(exc).__name__}: {exc}"}
finally:
if path and os.path.exists(path):
os.unlink(path)
runpod.serverless.start({"handler": handler})
+2
View File
@@ -0,0 +1,2 @@
runpod==1.7.9
faster-whisper==1.0.3
File diff suppressed because one or more lines are too long