Create monorepo from known-good production state

bgbye/server/ormbg.py

@@ -0,0 +1,70 @@
+# /app/server/ormbg.py
+from __future__ import annotations
+import io
+import os
+from typing import Optional, Union
+
+from rembg import remove as rembg_remove, new_session
+from PIL import Image
+import onnxruntime as ort
+
+# Choose GPU if available, else CPU
+def _default_providers() -> list[str]:
+    providers = ort.get_available_providers()
+    return ["CUDAExecutionProvider", "CPUExecutionProvider"] if "CUDAExecutionProvider" in providers else ["CPUExecutionProvider"]
+
+class ORMBGProcessor:
+    """
+    Minimal adapter used by server.py.
+    Provides:
+      - .warmup()
+      - .process_bytes(input_bytes) -> bytes (PNG with alpha)
+      - .process_pil(img: PIL.Image.Image) -> PIL.Image.Image
+      - .close()  (no-op; for API symmetry)
+    """
+    def __init__(self, model_name: str = "isnet-general-use", providers: Optional[list[str]] = None):
+        # Some rembg internals pick providers automatically; onnxruntime-gpu presence
+        # will make CUDA provider available. We still compute a list so we can log/verify.
+        self.providers = providers or _default_providers()
+        # rembg's API chooses the best available EP under the hood if ORT-GPU is present.
+        # Just ensure the model name exists; common: isnet-general-use, u2net, u2netp, u2net_human_seg
+        self.model_name = model_name
+        self.session = new_session(self.model_name)
+
+    def warmup(self) -> None:
+        # Tiny 1x1 transparent PNG to trigger model load
+        img = Image.new("RGBA", (1, 1), (0, 0, 0, 0))
+        _ = self.process_pil(img)
+
+    def process_bytes(self, data: bytes, return_png: bool = True) -> bytes:
+        """
+        Accept raw image bytes (jpg/png/etc), return PNG bytes with alpha.
+        """
+        out = rembg_remove(data, session=self.session)
+        if return_png:
+            # rembg already returns PNG bytes by default; if it's bytes, just pass through
+            if isinstance(out, (bytes, bytearray)):
+                return bytes(out)
+            # if it's a PIL image, convert to PNG bytes
+            if isinstance(out, Image.Image):
+                buf = io.BytesIO()
+                out.save(buf, format="PNG")
+                return buf.getvalue()
+        # fallback: ensure bytes
+        if isinstance(out, Image.Image):
+            buf = io.BytesIO()
+            out.save(buf, format="PNG")
+            return buf.getvalue()
+        return bytes(out)
+
+    def process_pil(self, img: Image.Image) -> Image.Image:
+        out = rembg_remove(img, session=self.session)
+        if isinstance(out, Image.Image):
+            return out
+        # if rembg returned bytes, decode back to PIL
+        return Image.open(io.BytesIO(out)).convert("RGBA")
+
+    def close(self) -> None:
+        # Nothing to dispose explicitly; keep for API compatibility
+        pass
+