Refactor code structure for improved readability and maintainability
This commit is contained in:
@@ -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"]
|
||||
@@ -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})
|
||||
@@ -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
|
||||
Reference in New Issue
Block a user