Refactor code structure for improved readability and maintainability

workers/cpu-tasks/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})