Create monorepo from known-good production state

bgbye/server/__init__.py

@@ -0,0 +1,3 @@
+# server/__init__.py
+__all__ = []
+

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
+

bgbye/server/ormbg/__init__.py

@@ -0,0 +1,2 @@
+from .ormbg import ORMBG
+from .ormbg_processor import ORMBGProcessor

bgbye/server/ormbg/ormbg.py

@@ -0,0 +1,484 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+# https://github.com/xuebinqin/DIS/blob/main/IS-Net/models/isnet.py
+
+
+class REBNCONV(nn.Module):
+    def __init__(self, in_ch=3, out_ch=3, dirate=1, stride=1):
+        super(REBNCONV, self).__init__()
+
+        self.conv_s1 = nn.Conv2d(
+            in_ch, out_ch, 3, padding=1 * dirate, dilation=1 * dirate, stride=stride
+        )
+        self.bn_s1 = nn.BatchNorm2d(out_ch)
+        self.relu_s1 = nn.ReLU(inplace=True)
+
+    def forward(self, x):
+
+        hx = x
+        xout = self.relu_s1(self.bn_s1(self.conv_s1(hx)))
+
+        return xout
+
+
+## upsample tensor 'src' to have the same spatial size with tensor 'tar'
+def _upsample_like(src, tar):
+
+    src = F.interpolate(src, size=tar.shape[2:], mode="bilinear")
+
+    return src
+
+
+### RSU-7 ###
+class RSU7(nn.Module):
+
+    def __init__(self, in_ch=3, mid_ch=12, out_ch=3, img_size=512):
+        super(RSU7, self).__init__()
+
+        self.in_ch = in_ch
+        self.mid_ch = mid_ch
+        self.out_ch = out_ch
+
+        self.rebnconvin = REBNCONV(in_ch, out_ch, dirate=1)  ## 1 -> 1/2
+
+        self.rebnconv1 = REBNCONV(out_ch, mid_ch, dirate=1)
+        self.pool1 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
+
+        self.rebnconv2 = REBNCONV(mid_ch, mid_ch, dirate=1)
+        self.pool2 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
+
+        self.rebnconv3 = REBNCONV(mid_ch, mid_ch, dirate=1)
+        self.pool3 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
+
+        self.rebnconv4 = REBNCONV(mid_ch, mid_ch, dirate=1)
+        self.pool4 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
+
+        self.rebnconv5 = REBNCONV(mid_ch, mid_ch, dirate=1)
+        self.pool5 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
+
+        self.rebnconv6 = REBNCONV(mid_ch, mid_ch, dirate=1)
+
+        self.rebnconv7 = REBNCONV(mid_ch, mid_ch, dirate=2)
+
+        self.rebnconv6d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)
+        self.rebnconv5d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)
+        self.rebnconv4d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)
+        self.rebnconv3d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)
+        self.rebnconv2d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)
+        self.rebnconv1d = REBNCONV(mid_ch * 2, out_ch, dirate=1)
+
+    def forward(self, x):
+        b, c, h, w = x.shape
+
+        hx = x
+        hxin = self.rebnconvin(hx)
+
+        hx1 = self.rebnconv1(hxin)
+        hx = self.pool1(hx1)
+
+        hx2 = self.rebnconv2(hx)
+        hx = self.pool2(hx2)
+
+        hx3 = self.rebnconv3(hx)
+        hx = self.pool3(hx3)
+
+        hx4 = self.rebnconv4(hx)
+        hx = self.pool4(hx4)
+
+        hx5 = self.rebnconv5(hx)
+        hx = self.pool5(hx5)
+
+        hx6 = self.rebnconv6(hx)
+
+        hx7 = self.rebnconv7(hx6)
+
+        hx6d = self.rebnconv6d(torch.cat((hx7, hx6), 1))
+        hx6dup = _upsample_like(hx6d, hx5)
+
+        hx5d = self.rebnconv5d(torch.cat((hx6dup, hx5), 1))
+        hx5dup = _upsample_like(hx5d, hx4)
+
+        hx4d = self.rebnconv4d(torch.cat((hx5dup, hx4), 1))
+        hx4dup = _upsample_like(hx4d, hx3)
+
+        hx3d = self.rebnconv3d(torch.cat((hx4dup, hx3), 1))
+        hx3dup = _upsample_like(hx3d, hx2)
+
+        hx2d = self.rebnconv2d(torch.cat((hx3dup, hx2), 1))
+        hx2dup = _upsample_like(hx2d, hx1)
+
+        hx1d = self.rebnconv1d(torch.cat((hx2dup, hx1), 1))
+
+        return hx1d + hxin
+
+
+### RSU-6 ###
+class RSU6(nn.Module):
+
+    def __init__(self, in_ch=3, mid_ch=12, out_ch=3):
+        super(RSU6, self).__init__()
+
+        self.rebnconvin = REBNCONV(in_ch, out_ch, dirate=1)
+
+        self.rebnconv1 = REBNCONV(out_ch, mid_ch, dirate=1)
+        self.pool1 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
+
+        self.rebnconv2 = REBNCONV(mid_ch, mid_ch, dirate=1)
+        self.pool2 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
+
+        self.rebnconv3 = REBNCONV(mid_ch, mid_ch, dirate=1)
+        self.pool3 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
+
+        self.rebnconv4 = REBNCONV(mid_ch, mid_ch, dirate=1)
+        self.pool4 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
+
+        self.rebnconv5 = REBNCONV(mid_ch, mid_ch, dirate=1)
+
+        self.rebnconv6 = REBNCONV(mid_ch, mid_ch, dirate=2)
+
+        self.rebnconv5d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)
+        self.rebnconv4d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)
+        self.rebnconv3d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)
+        self.rebnconv2d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)
+        self.rebnconv1d = REBNCONV(mid_ch * 2, out_ch, dirate=1)
+
+    def forward(self, x):
+
+        hx = x
+
+        hxin = self.rebnconvin(hx)
+
+        hx1 = self.rebnconv1(hxin)
+        hx = self.pool1(hx1)
+
+        hx2 = self.rebnconv2(hx)
+        hx = self.pool2(hx2)
+
+        hx3 = self.rebnconv3(hx)
+        hx = self.pool3(hx3)
+
+        hx4 = self.rebnconv4(hx)
+        hx = self.pool4(hx4)
+
+        hx5 = self.rebnconv5(hx)
+
+        hx6 = self.rebnconv6(hx5)
+
+        hx5d = self.rebnconv5d(torch.cat((hx6, hx5), 1))
+        hx5dup = _upsample_like(hx5d, hx4)
+
+        hx4d = self.rebnconv4d(torch.cat((hx5dup, hx4), 1))
+        hx4dup = _upsample_like(hx4d, hx3)
+
+        hx3d = self.rebnconv3d(torch.cat((hx4dup, hx3), 1))
+        hx3dup = _upsample_like(hx3d, hx2)
+
+        hx2d = self.rebnconv2d(torch.cat((hx3dup, hx2), 1))
+        hx2dup = _upsample_like(hx2d, hx1)
+
+        hx1d = self.rebnconv1d(torch.cat((hx2dup, hx1), 1))
+
+        return hx1d + hxin
+
+
+### RSU-5 ###
+class RSU5(nn.Module):
+
+    def __init__(self, in_ch=3, mid_ch=12, out_ch=3):
+        super(RSU5, self).__init__()
+
+        self.rebnconvin = REBNCONV(in_ch, out_ch, dirate=1)
+
+        self.rebnconv1 = REBNCONV(out_ch, mid_ch, dirate=1)
+        self.pool1 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
+
+        self.rebnconv2 = REBNCONV(mid_ch, mid_ch, dirate=1)
+        self.pool2 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
+
+        self.rebnconv3 = REBNCONV(mid_ch, mid_ch, dirate=1)
+        self.pool3 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
+
+        self.rebnconv4 = REBNCONV(mid_ch, mid_ch, dirate=1)
+
+        self.rebnconv5 = REBNCONV(mid_ch, mid_ch, dirate=2)
+
+        self.rebnconv4d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)
+        self.rebnconv3d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)
+        self.rebnconv2d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)
+        self.rebnconv1d = REBNCONV(mid_ch * 2, out_ch, dirate=1)
+
+    def forward(self, x):
+
+        hx = x
+
+        hxin = self.rebnconvin(hx)
+
+        hx1 = self.rebnconv1(hxin)
+        hx = self.pool1(hx1)
+
+        hx2 = self.rebnconv2(hx)
+        hx = self.pool2(hx2)
+
+        hx3 = self.rebnconv3(hx)
+        hx = self.pool3(hx3)
+
+        hx4 = self.rebnconv4(hx)
+
+        hx5 = self.rebnconv5(hx4)
+
+        hx4d = self.rebnconv4d(torch.cat((hx5, hx4), 1))
+        hx4dup = _upsample_like(hx4d, hx3)
+
+        hx3d = self.rebnconv3d(torch.cat((hx4dup, hx3), 1))
+        hx3dup = _upsample_like(hx3d, hx2)
+
+        hx2d = self.rebnconv2d(torch.cat((hx3dup, hx2), 1))
+        hx2dup = _upsample_like(hx2d, hx1)
+
+        hx1d = self.rebnconv1d(torch.cat((hx2dup, hx1), 1))
+
+        return hx1d + hxin
+
+
+### RSU-4 ###
+class RSU4(nn.Module):
+
+    def __init__(self, in_ch=3, mid_ch=12, out_ch=3):
+        super(RSU4, self).__init__()
+
+        self.rebnconvin = REBNCONV(in_ch, out_ch, dirate=1)
+
+        self.rebnconv1 = REBNCONV(out_ch, mid_ch, dirate=1)
+        self.pool1 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
+
+        self.rebnconv2 = REBNCONV(mid_ch, mid_ch, dirate=1)
+        self.pool2 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
+
+        self.rebnconv3 = REBNCONV(mid_ch, mid_ch, dirate=1)
+
+        self.rebnconv4 = REBNCONV(mid_ch, mid_ch, dirate=2)
+
+        self.rebnconv3d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)
+        self.rebnconv2d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)
+        self.rebnconv1d = REBNCONV(mid_ch * 2, out_ch, dirate=1)
+
+    def forward(self, x):
+
+        hx = x
+
+        hxin = self.rebnconvin(hx)
+
+        hx1 = self.rebnconv1(hxin)
+        hx = self.pool1(hx1)
+
+        hx2 = self.rebnconv2(hx)
+        hx = self.pool2(hx2)
+
+        hx3 = self.rebnconv3(hx)
+
+        hx4 = self.rebnconv4(hx3)
+
+        hx3d = self.rebnconv3d(torch.cat((hx4, hx3), 1))
+        hx3dup = _upsample_like(hx3d, hx2)
+
+        hx2d = self.rebnconv2d(torch.cat((hx3dup, hx2), 1))
+        hx2dup = _upsample_like(hx2d, hx1)
+
+        hx1d = self.rebnconv1d(torch.cat((hx2dup, hx1), 1))
+
+        return hx1d + hxin
+
+
+### RSU-4F ###
+class RSU4F(nn.Module):
+
+    def __init__(self, in_ch=3, mid_ch=12, out_ch=3):
+        super(RSU4F, self).__init__()
+
+        self.rebnconvin = REBNCONV(in_ch, out_ch, dirate=1)
+
+        self.rebnconv1 = REBNCONV(out_ch, mid_ch, dirate=1)
+        self.rebnconv2 = REBNCONV(mid_ch, mid_ch, dirate=2)
+        self.rebnconv3 = REBNCONV(mid_ch, mid_ch, dirate=4)
+
+        self.rebnconv4 = REBNCONV(mid_ch, mid_ch, dirate=8)
+
+        self.rebnconv3d = REBNCONV(mid_ch * 2, mid_ch, dirate=4)
+        self.rebnconv2d = REBNCONV(mid_ch * 2, mid_ch, dirate=2)
+        self.rebnconv1d = REBNCONV(mid_ch * 2, out_ch, dirate=1)
+
+    def forward(self, x):
+
+        hx = x
+
+        hxin = self.rebnconvin(hx)
+
+        hx1 = self.rebnconv1(hxin)
+        hx2 = self.rebnconv2(hx1)
+        hx3 = self.rebnconv3(hx2)
+
+        hx4 = self.rebnconv4(hx3)
+
+        hx3d = self.rebnconv3d(torch.cat((hx4, hx3), 1))
+        hx2d = self.rebnconv2d(torch.cat((hx3d, hx2), 1))
+        hx1d = self.rebnconv1d(torch.cat((hx2d, hx1), 1))
+
+        return hx1d + hxin
+
+
+class myrebnconv(nn.Module):
+    def __init__(
+        self,
+        in_ch=3,
+        out_ch=1,
+        kernel_size=3,
+        stride=1,
+        padding=1,
+        dilation=1,
+        groups=1,
+    ):
+        super(myrebnconv, self).__init__()
+
+        self.conv = nn.Conv2d(
+            in_ch,
+            out_ch,
+            kernel_size=kernel_size,
+            stride=stride,
+            padding=padding,
+            dilation=dilation,
+            groups=groups,
+        )
+        self.bn = nn.BatchNorm2d(out_ch)
+        self.rl = nn.ReLU(inplace=True)
+
+    def forward(self, x):
+        return self.rl(self.bn(self.conv(x)))
+
+
+bce_loss = nn.BCELoss(size_average=True)
+
+
+class ORMBG(nn.Module):
+
+    def __init__(self, in_ch=3, out_ch=1):
+        super(ORMBG, self).__init__()
+
+        self.conv_in = nn.Conv2d(in_ch, 64, 3, stride=2, padding=1)
+        self.pool_in = nn.MaxPool2d(2, stride=2, ceil_mode=True)
+
+        self.stage1 = RSU7(64, 32, 64)
+        self.pool12 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
+
+        self.stage2 = RSU6(64, 32, 128)
+        self.pool23 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
+
+        self.stage3 = RSU5(128, 64, 256)
+        self.pool34 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
+
+        self.stage4 = RSU4(256, 128, 512)
+        self.pool45 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
+
+        self.stage5 = RSU4F(512, 256, 512)
+        self.pool56 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
+
+        self.stage6 = RSU4F(512, 256, 512)
+
+        # decoder
+        self.stage5d = RSU4F(1024, 256, 512)
+        self.stage4d = RSU4(1024, 128, 256)
+        self.stage3d = RSU5(512, 64, 128)
+        self.stage2d = RSU6(256, 32, 64)
+        self.stage1d = RSU7(128, 16, 64)
+
+        self.side1 = nn.Conv2d(64, out_ch, 3, padding=1)
+        self.side2 = nn.Conv2d(64, out_ch, 3, padding=1)
+        self.side3 = nn.Conv2d(128, out_ch, 3, padding=1)
+        self.side4 = nn.Conv2d(256, out_ch, 3, padding=1)
+        self.side5 = nn.Conv2d(512, out_ch, 3, padding=1)
+        self.side6 = nn.Conv2d(512, out_ch, 3, padding=1)
+
+        # self.outconv = nn.Conv2d(6*out_ch,out_ch,1)
+
+    def compute_loss(self, predictions, ground_truth):
+        loss0, loss = 0.0, 0.0
+        for i in range(0, len(predictions)):
+            loss = loss + bce_loss(predictions[i], ground_truth)
+            if i == 0:
+                loss0 = loss
+        return loss0, loss
+
+    def forward(self, x):
+
+        hx = x
+
+        hxin = self.conv_in(hx)
+        # hx = self.pool_in(hxin)
+
+        # stage 1
+        hx1 = self.stage1(hxin)
+        hx = self.pool12(hx1)
+
+        # stage 2
+        hx2 = self.stage2(hx)
+        hx = self.pool23(hx2)
+
+        # stage 3
+        hx3 = self.stage3(hx)
+        hx = self.pool34(hx3)
+
+        # stage 4
+        hx4 = self.stage4(hx)
+        hx = self.pool45(hx4)
+
+        # stage 5
+        hx5 = self.stage5(hx)
+        hx = self.pool56(hx5)
+
+        # stage 6
+        hx6 = self.stage6(hx)
+        hx6up = _upsample_like(hx6, hx5)
+
+        # -------------------- decoder --------------------
+        hx5d = self.stage5d(torch.cat((hx6up, hx5), 1))
+        hx5dup = _upsample_like(hx5d, hx4)
+
+        hx4d = self.stage4d(torch.cat((hx5dup, hx4), 1))
+        hx4dup = _upsample_like(hx4d, hx3)
+
+        hx3d = self.stage3d(torch.cat((hx4dup, hx3), 1))
+        hx3dup = _upsample_like(hx3d, hx2)
+
+        hx2d = self.stage2d(torch.cat((hx3dup, hx2), 1))
+        hx2dup = _upsample_like(hx2d, hx1)
+
+        hx1d = self.stage1d(torch.cat((hx2dup, hx1), 1))
+
+        # side output
+        d1 = self.side1(hx1d)
+        d1 = _upsample_like(d1, x)
+
+        d2 = self.side2(hx2d)
+        d2 = _upsample_like(d2, x)
+
+        d3 = self.side3(hx3d)
+        d3 = _upsample_like(d3, x)
+
+        d4 = self.side4(hx4d)
+        d4 = _upsample_like(d4, x)
+
+        d5 = self.side5(hx5d)
+        d5 = _upsample_like(d5, x)
+
+        d6 = self.side6(hx6)
+        d6 = _upsample_like(d6, x)
+
+        return [
+            F.sigmoid(d1),
+            F.sigmoid(d2),
+            F.sigmoid(d3),
+            F.sigmoid(d4),
+            F.sigmoid(d5),
+            F.sigmoid(d6),
+        ], [hx1d, hx2d, hx3d, hx4d, hx5d, hx6]

bgbye/server/ormbg/ormbg_processor.py

@@ -0,0 +1,45 @@
+import torch
+import torch.nn.functional as F
+import numpy as np
+from PIL import Image
+from .ormbg import ORMBG
+
+class ORMBGProcessor:
+    def __init__(self, model_path):
+        self.device = torch.device("cpu")
+        self.net = ORMBG()
+        self.net.load_state_dict(torch.load(model_path, map_location="cpu"))
+        self.net.eval()
+
+    def to(self, device):
+        self.device = torch.device(device)
+        self.net.to(self.device)
+
+    def process_image(self, image):
+        # Ensure image is in RGB mode
+        if image.mode != "RGB":
+            image = image.convert("RGB")
+
+        # Preprocess the image
+        w, h = image.size
+        image = image.resize((1024, 1024), Image.BILINEAR)
+        im_np = np.array(image)
+        im_tensor = torch.tensor(im_np, dtype=torch.float32).permute(2, 0, 1).unsqueeze(0)
+        im_tensor = torch.divide(im_tensor, 255.0).to(self.device)
+
+        # Inference
+        with torch.no_grad():
+            result = self.net(im_tensor)
+
+        # Post-process
+        result = result[0][0]  # Take the first element of the output list and the first channel
+        result = F.interpolate(result, size=(h, w), mode="bilinear")
+        result = result.squeeze()
+        result = (result - result.min()) / (result.max() - result.min())
+        
+        # Create mask and apply to original image
+        mask = Image.fromarray((result.cpu().numpy() * 255).astype(np.uint8))
+        new_im = Image.new("RGBA", (w, h), (0, 0, 0, 0))
+        new_im.paste(image.resize((w, h), Image.BILINEAR), mask=mask)
+
+        return new_im

bgbye/server/run.sh

@@ -0,0 +1,20 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+MODEL_PATH="${ORMBG_MODEL_PATH:-/root/.ormbg/ormbg.pth}"
+MODEL_DIR="$(dirname "$MODEL_PATH")"
+
+if [ ! -f "$MODEL_PATH" ]; then
+  echo "[bgbye] ORMBG model not found at $MODEL_PATH - running setup.sh to fetch it..."
+  mkdir -p "$MODEL_DIR"
+  if [ -x /app/server/setup.sh ]; then
+    bash /app/server/setup.sh
+  else
+    echo "[bgbye] ERROR: /app/server/setup.sh not found or not executable" >&2
+    exit 1
+  fi
+fi
+
+echo "[bgbye] Starting API..."
+exec uvicorn server.server:app --host 0.0.0.0 --port 7001 --no-server-header
+

bgbye/server/server.py

@@ -0,0 +1,513 @@
+from fastapi import FastAPI, UploadFile, File, Response, Form, HTTPException, BackgroundTasks
+from fastapi.middleware.cors import CORSMiddleware
+from fastapi.responses import FileResponse, JSONResponse
+from PIL import Image
+import io
+import shutil
+from rembg import remove as rembg_remove, new_session
+import time
+import numpy as np
+import tempfile
+import uuid
+import os
+import subprocess
+from transformers import pipeline
+from transparent_background import Remover
+import logging
+import asyncio
+from datetime import datetime, timedelta
+import torch
+from .ormbg import ORMBGProcessor
+from typing import Dict
+from contextlib import contextmanager
+
+from carvekit.ml.files.models_loc import download_all
+
+download_all()
+
+from carvekit.ml.wrap.u2net import U2NET
+from carvekit.ml.wrap.basnet import BASNET
+from carvekit.ml.wrap.fba_matting import FBAMatting
+from carvekit.ml.wrap.deeplab_v3 import DeepLabV3
+from carvekit.ml.wrap.tracer_b7 import TracerUniversalB7
+from carvekit.api.interface import Interface
+from carvekit.pipelines.postprocessing import MattingMethod
+from carvekit.pipelines.preprocessing import PreprocessingStub
+from carvekit.trimap.generator import TrimapGenerator
+
+
+# Set up logging
+logging.basicConfig(level=logging.DEBUG)
+logger = logging.getLogger(__name__)
+
+# Add ORMBG model initialization
+ormbg_model_path = os.path.expanduser("~/.ormbg/ormbg.pth")
+try:
+    ormbg_processor = ORMBGProcessor(ormbg_model_path)
+    if torch.cuda.is_available():
+        ormbg_processor.to("cuda")
+    else:
+        ormbg_processor.to("cpu")
+except FileNotFoundError:
+    logger.error(f"ORMBG model file not found: {ormbg_model_path}")
+    print("Error: ORMBG model file not found. Please run 'npm run setup-server' to download it.")
+    exit(1)
+
+app = FastAPI()
+
+# Create temp_videos folder if it doesn't exist
+TEMP_VIDEOS_DIR = "temp_videos"
+os.makedirs(TEMP_VIDEOS_DIR, exist_ok=True)
+
+# Create a frames directory within temp_videos
+FRAMES_DIR = os.path.join(TEMP_VIDEOS_DIR, "frames")
+os.makedirs(FRAMES_DIR, exist_ok=True)
+
+# Add a dictionary to store processing status
+processing_status = {}
+
+# Add CORS middleware
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+async def cleanup_old_videos():
+    while True:
+        current_time = datetime.now()
+        for item in os.listdir(TEMP_VIDEOS_DIR):
+            item_path = os.path.join(TEMP_VIDEOS_DIR, item)
+            item_modified = datetime.fromtimestamp(os.path.getmtime(item_path))
+            if current_time - item_modified > timedelta(minutes=10):
+                if os.path.isfile(item_path):
+                    os.remove(item_path)
+                    logger.info(f"Removed old file: {item_path}")
+                elif os.path.isdir(item_path):
+                    shutil.rmtree(item_path)
+                    logger.info(f"Removed old directory: {item_path}")
+        await asyncio.sleep(600)  # Run every 10 minutes
+
+# Pre-load all models
+bria_model = pipeline("image-segmentation", model="briaai/RMBG-1.4", trust_remote_code=True, device="cpu")
+inspyrenet_model = Remover()
+inspyrenet_model.model.cpu()
+rembg_models = {
+    'u2net': new_session('u2net'),
+    'u2net_human_seg': new_session('u2net_human_seg'),
+    'isnet-general-use': new_session('isnet-general-use'),
+    'isnet-anime': new_session('isnet-anime')
+}
+
+# Initialize Carvekit models
+def initialize_carvekit_model(seg_pipe_class, device='cuda'):
+    model = Interface(
+        pre_pipe=PreprocessingStub(),
+        post_pipe=MattingMethod(
+            matting_module=FBAMatting(device=device, input_tensor_size=2048, batch_size=1),
+            trimap_generator=TrimapGenerator(),
+            device=device
+        ),
+        seg_pipe=seg_pipe_class(device=device, batch_size=1)
+    )
+    model.segmentation_pipeline.to('cpu')
+    return model
+
+carvekit_models = {
+    'u2net': initialize_carvekit_model(U2NET),
+    'tracer': initialize_carvekit_model(TracerUniversalB7),
+    'basnet': initialize_carvekit_model(BASNET),
+    'deeplab': initialize_carvekit_model(DeepLabV3)
+}
+
+# ---- Method registry + aliases (NEW) -----------------------------------------
+CARVEKIT = list(carvekit_models.keys())          # ['u2net', 'tracer', 'basnet', 'deeplab']
+REMBG    = ['u2net_human_seg', 'isnet-general-use', 'isnet-anime']
+EXTRA    = ['ormbg', 'bria', 'inspyrenet']
+ALLOWED_METHODS = sorted(set(CARVEKIT + REMBG + EXTRA))
+
+ALIASES = {
+    'deeplabv3': 'deeplab',
+    'tracer-b7': 'tracer',
+    'tracer_b7': 'tracer',
+    'u2net_human': 'u2net_human_seg',
+    'isnet_general_use': 'isnet-general-use',
+    'bria14': 'bria',
+    # (intentionally no alias for 'rembg'/'open-rmbg')
+}
+
+@app.get("/methods")
+def list_methods():
+    return {"methods": ALLOWED_METHODS}
+
+@app.get("/health")
+def health():
+    return JSONResponse({"ok": True})
+# -----------------------------------------------------------------------------
+
+# Ensure GPU memory is cleared after initialization
+torch.cuda.empty_cache()
+
+def process_with_bria(image):
+    result = bria_model(image, return_mask=True)
+    mask = result
+    if not isinstance(mask, Image.Image):
+        mask = Image.fromarray((mask * 255).astype('uint8'))
+    no_bg_image = Image.new("RGBA", image.size, (0, 0, 0, 0))
+    no_bg_image.paste(image, mask=mask)
+    return no_bg_image
+
+def process_with_ormbg(image):
+    result = ormbg_processor.process_image(image)
+    return result
+
+def process_with_inspyrenet(image):
+    return inspyrenet_model.process(image, type='rgba')
+
+def process_with_rembg(image, model='u2net'):
+    return rembg_remove(image, session=rembg_models[model])
+
+def process_with_carvekit(image, model='u2net'):
+    # Initialize segmentation network based on model input
+    if model == 'u2net':
+        seg_net = U2NET(device='cuda', batch_size=1)
+    elif model == 'tracer':
+        seg_net = TracerUniversalB7(device='cuda', batch_size=1)
+    elif model == 'basnet':
+        seg_net = BASNET(device='cuda', batch_size=1)
+    elif model == 'deeplab':
+        seg_net = DeepLabV3(device='cuda', batch_size=1)
+    else:
+        raise ValueError("Unsupported model type")
+
+    # Setup the post-processing components
+    fba = FBAMatting(device='cuda', input_tensor_size=2048, batch_size=1)
+    trimap = TrimapGenerator()
+    preprocessing = PreprocessingStub()
+    postprocessing = MattingMethod(matting_module=fba, trimap_generator=trimap, device='cuda')
+
+    interface = Interface(pre_pipe=preprocessing, post_pipe=postprocessing, seg_pipe=seg_net)
+    processed_image = interface([image])[0]
+    
+    return processed_image
+
+@contextmanager
+def inspyrenet_video_model_context():
+    try:
+        model = Remover()
+        model.model.cuda()
+        yield model
+    finally:
+        model.model.cpu()
+        del model
+        torch.cuda.empty_cache()
+
+@contextmanager
+def carvekit_video_model_context(model_name):
+    try:
+        if model_name == 'u2net':
+            seg_net = U2NET(device='cuda', batch_size=1)
+        elif model_name == 'tracer':
+            seg_net = TracerUniversalB7(device='cuda', batch_size=1)
+        elif model_name == 'basnet':
+            seg_net = BASNET(device='cuda', batch_size=1)
+        elif model_name == 'deeplab':
+            seg_net = DeepLabV3(device='cuda', batch_size=1)
+        else:
+            raise ValueError("Unsupported model type")
+
+        fba = FBAMatting(device='cuda', input_tensor_size=2048, batch_size=1)
+        trimap = TrimapGenerator()
+        preprocessing = PreprocessingStub()
+        postprocessing = MattingMethod(matting_module=fba, trimap_generator=trimap, device='cuda')
+
+        interface = Interface(pre_pipe=preprocessing, post_pipe=postprocessing, seg_pipe=seg_net)
+        yield interface
+    finally:
+        del seg_net, fba, trimap, preprocessing, postprocessing, interface
+        torch.cuda.empty_cache()
+
+
+# Create a global lock for GPU operations
+gpu_lock = asyncio.Lock()
+
+@app.post("/remove_background/")
+async def remove_background(file: UploadFile = File(...), method: str = Form(...)):
+    try:
+        # normalize aliases
+        method = ALIASES.get(method, method)
+
+        # early reject unknown methods with 400, not 500
+        if method not in ALLOWED_METHODS:
+            raise HTTPException(status_code=400, detail="Invalid method")
+
+        image_data = await file.read()
+        image = Image.open(io.BytesIO(image_data)).convert('RGB')
+        
+        start_time = time.time()
+
+        async def process_image():
+            if method == 'bria':
+                return await asyncio.to_thread(process_with_bria, image)
+            elif method == 'inspyrenet':
+                async with gpu_lock:
+                    try:
+                        inspyrenet_model.model.to('cuda')
+                        result = await asyncio.to_thread(inspyrenet_model.process, image, type='rgba')
+                    finally:
+                        inspyrenet_model.model.to('cpu')
+                    return result
+            elif method in ['u2net_human_seg', 'isnet-general-use', 'isnet-anime']:
+                return await asyncio.to_thread(process_with_rembg, image, model=method)
+            elif method == 'ormbg':
+                return await asyncio.to_thread(process_with_ormbg, image)
+            elif method in ['u2net', 'tracer', 'basnet', 'deeplab']:
+                async with gpu_lock:
+                    try:
+                        carvekit_models[method].segmentation_pipeline.to('cuda')
+                        result = await asyncio.to_thread(carvekit_models[method], [image])
+                    finally:
+                        carvekit_models[method].segmentation_pipeline.to('cpu')
+                    return result[0]
+            else:
+                raise HTTPException(status_code=400, detail="Invalid method")
+
+        no_bg_image = await process_image()
+        
+        process_time = time.time() - start_time
+        print(f"Background removal time ({method}): {process_time:.2f} seconds")
+        
+        async with gpu_lock:
+            torch.cuda.empty_cache()
+        
+        with io.BytesIO() as output:
+            no_bg_image.save(output, format="PNG")
+            content = output.getvalue()
+
+        return Response(content=content, media_type="image/png")
+
+    except HTTPException as he:
+        # let 4xx bubble as-is
+        raise he
+    except Exception as e:
+        print(str(e))
+        raise HTTPException(status_code=500, detail=str(e))
+
+async def process_frame(frame_path, method):
+    img = Image.open(frame_path).convert('RGB')
+    
+    if method == 'bria':
+        processed_frame = await asyncio.to_thread(process_with_bria, img)
+    elif method in ['u2net_human_seg', 'isnet-general-use', 'isnet-anime']:
+        processed_frame = await asyncio.to_thread(process_with_rembg, img, model=method)
+    elif method == 'ormbg':
+        processed_frame = await asyncio.to_thread(process_with_ormbg, img)
+    else:
+        raise ValueError("Invalid method")
+    
+    return processed_frame
+
+async def process_video(video_path, method, video_id):
+    try:
+        processing_status[video_id] = {'status': 'processing', 'progress': 0, 'message': 'Initializing'}
+        
+        logger.info(f"Starting video processing: {video_path}")
+        logger.info(f"Method: {method}")
+        logger.info(f"Video ID: {video_id}")
+
+
+        # Check video frame count
+        frame_count_command = ['ffmpeg.ffprobe', '-v', 'error', '-select_streams', 'v:0', '-count_packets', 
+                               '-show_entries', 'stream=nb_read_packets', '-of', 'csv=p=0', video_path]
+        process = await asyncio.create_subprocess_exec(
+            *frame_count_command,
+            stdout=asyncio.subprocess.PIPE,
+            stderr=asyncio.subprocess.PIPE
+        )
+        stdout, stderr = await process.communicate()
+        
+        if process.returncode != 0:
+            logger.error(f"Error counting frames: {stderr.decode()}")
+            processing_status[video_id] = {'status': 'error', 'message': 'Error counting frames'}
+            return
+
+        frame_count = int(stdout.decode().strip())
+        logger.info(f"Video frame count: {frame_count}")
+
+        #DISABLED VIDEO LENGTH LIMIT
+        #if frame_count > 250:
+        #    logger.warning(f"Video too long: {frame_count} frames")
+        #    processing_status[video_id] = {'status': 'error', 'message': 'Video too long (max 250 frames)'}
+        #    return
+
+        # Create a unique directory for this video's frames
+        frames_dir = os.path.join(FRAMES_DIR, video_id)
+        os.makedirs(frames_dir, exist_ok=True)
+        logger.info(f"Created frames directory: {frames_dir}")
+
+        # Extract frames from video
+        processing_status[video_id] = {'status': 'processing', 'progress': 0, 'message': 'Extracting frames'}
+        extract_command = ['ffmpeg', '-i', video_path, f'{frames_dir}/frame_%05d.png']
+        logger.info(f"Executing frame extraction command: {' '.join(extract_command)}")
+        process = await asyncio.create_subprocess_exec(
+            *extract_command,
+            stdout=asyncio.subprocess.PIPE,
+            stderr=asyncio.subprocess.PIPE
+        )
+        stdout, stderr = await process.communicate()
+        
+        if process.returncode != 0:
+            logger.error(f"Error extracting frames: {stderr.decode()}")
+            processing_status[video_id] = {'status': 'error', 'message': 'Error extracting frames'}
+            return
+
+        # Process frames
+        processing_status[video_id] = {'status': 'processing', 'progress': 0, 'message': 'Removing background'}
+        frame_files = sorted([f for f in os.listdir(frames_dir) if f.endswith('.png')])
+        total_frames = len(frame_files)
+        logger.info(f"Number of extracted frames: {total_frames}")
+
+        if total_frames == 0:
+            logger.error("No frames were extracted from the video")
+            processing_status[video_id] = {'status': 'error', 'message': 'No frames were extracted from the video'}
+            return
+
+        # Initialize the model once, outside the batch processing loop
+        if method == 'inspyrenet':
+            print("start init")
+            model_context = inspyrenet_video_model_context()
+            print("start enter")
+            model = model_context.__enter__()
+            print("finish enter")
+        elif method in ['u2net', 'tracer', 'basnet', 'deeplab']:
+            model_context = carvekit_video_model_context(method)
+            model = model_context.__enter__()
+        else:
+            model = None  # For other methods that don't require a specific model
+
+        try:
+            async def process_frame_batch(start_idx, end_idx):
+                for i in range(start_idx, min(end_idx, total_frames)):
+                    frame_file = frame_files[i]
+                    frame_path = os.path.join(frames_dir, frame_file)
+                    img = Image.open(frame_path).convert('RGB')
+
+                    if method == 'inspyrenet':
+                        processed_frame = model.process(img, type='rgba')
+                    elif method in ['u2net', 'tracer', 'basnet', 'deeplab']:
+                        processed_frame = model([img])[0]
+                    else:
+                        processed_frame = await process_frame(frame_path, method)
+
+                    processed_frame.save(frame_path, format='PNG')
+                    progress = (i + 1) / total_frames * 100
+                    processing_status[video_id] = {'status': 'processing', 'progress': progress}
+
+            batch_size = 3
+            for i in range(0, total_frames, batch_size):
+                await process_frame_batch(i, i + batch_size)
+                await asyncio.sleep(0)  # Allow other tasks to run
+
+        finally:
+            # Ensure we clean up the model context
+            if method in ['inspyrenet', 'u2net', 'tracer', 'basnet', 'deeplab']:
+                model_context.__exit__(None, None, None)
+
+        # Create output video
+        processing_status[video_id] = {'status': 'processing', 'progress': 100, 'message': 'Encoding video'}
+        output_path = os.path.join(TEMP_VIDEOS_DIR, f"output_{video_id}.webm")
+        create_video_command = [
+            'ffmpeg',
+            '-framerate', '24',
+            '-i', f'{frames_dir}/frame_%05d.png',
+            '-c:v', 'libvpx-vp9',
+            '-pix_fmt', 'yuva420p',
+            '-lossless', '1',
+            output_path
+        ]
+        logger.info(f"Executing video creation command: {' '.join(create_video_command)}")
+        process = await asyncio.create_subprocess_exec(
+            *create_video_command,
+            stdout=asyncio.subprocess.PIPE,
+            stderr=asyncio.subprocess.PIPE
+        )
+        stdout, stderr = await process.communicate()
+        
+        if process.returncode != 0:
+            logger.error(f"Error creating output video: {stderr.decode()}")
+            processing_status[video_id] = {'status': 'error', 'message': 'Error creating output video'}
+            return
+
+        logger.info(f"Video processing completed. Output path: {output_path}")
+        processing_status[video_id] = {'status': 'completed', 'output_path': output_path}
+
+    except Exception as e:
+        logger.exception("Error in video processing")
+        processing_status[video_id] = {'status': 'error', 'message': str(e)}
+    finally:
+        torch.cuda.empty_cache()
+
+        # Clean up frames directory
+        for file in os.listdir(frames_dir):
+            os.remove(os.path.join(frames_dir, file))
+        os.rmdir(frames_dir)
+        logger.info(f"Cleaned up frames directory: {frames_dir}")
+
+@app.post("/remove_background_video/")
+async def remove_background_video(background_tasks: BackgroundTasks, file: UploadFile = File(...), method: str = Form(...)):
+    try:
+        logger.info(f"Starting video background removal with method: {method}")
+        
+        # Generate a unique filename for the uploaded video
+        video_id = str(uuid.uuid4())
+        filename = f"input_{video_id}.mp4"
+        file_path = os.path.join(TEMP_VIDEOS_DIR, filename)
+        
+        # Save uploaded video to the temp_videos folder
+        with open(file_path, "wb") as buffer:
+            content = await file.read()
+            buffer.write(content)
+
+        logger.info(f"Video file saved: {file_path}")
+        logger.info(f"File exists: {os.path.exists(file_path)}")
+        logger.info(f"File size: {os.path.getsize(file_path)} bytes")
+
+        if not os.path.exists(file_path):
+            raise HTTPException(status_code=500, detail=f"Failed to create video file: {file_path}")
+
+        # Start processing in the background
+        background_tasks.add_task(process_video, file_path, method, video_id)
+        
+        return {"video_id": video_id}
+
+    except Exception as e:
+        logger.exception(f"Error in video processing: {str(e)}")
+        raise HTTPException(status_code=500, detail=f"Error in video processing: {str(e)}")
+
+@app.get("/status/{video_id}")
+async def get_status(video_id: str):
+    if video_id not in processing_status:
+        raise HTTPException(status_code=404, detail="Video ID not found")
+    
+    status = processing_status[video_id]
+    
+    if status['status'] == 'completed':
+        output_path = status['output_path']
+        if not os.path.exists(output_path):
+            raise HTTPException(status_code=404, detail="Processed video file not found")
+        
+        return FileResponse(output_path, media_type="video/webm", filename=f"processed_video_{video_id}.webm")
+    
+    return status
+
+@app.on_event("startup")
+async def startup_event():
+    asyncio.create_task(cleanup_old_videos())
+    
+
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=9876)
+

bgbye/server/setup.sh

@@ -0,0 +1,31 @@
+#!/bin/bash
+
+# Update package list
+sudo apt update
+
+# Install system dependencies
+sudo apt install -y python3-pip python3-venv pngcrush libjpeg-turbo-progs jpegoptim \
+    libgl1-mesa-glx libglib2.0-0 libsm6 libxext6 libxrender-dev ffmpeg
+
+# Create and activate a virtual environment
+python3 -m venv venv
+source venv/bin/activate
+
+# Install Python dependencies
+pip install numpy==1.26.4 #WEIRD IMPORT ERROR WORKAROUND FOR REMBG
+pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121
+pip install fastapi uvicorn transformers pillow scikit-image transparent-background rembg opencv-python-headless python-multipart requests
+pip install carvekit #--extra-index-url https://download.pytorch.org/whl/cu121
+
+# Ensure stuff is in the PATH
+echo 'export PATH=$PATH:$HOME/.local/bin' >> ~/.bashrc
+source ~/.bashrc
+
+# Install additional requirements if there's a requirements.txt file
+if [ -f requirements.txt ]; then
+    pip install -r requirements.txt
+fi
+
+mkdir -p ~/.ormbg && echo "Downloading the ORMBG model..." && wget -O ~/.ormbg/ormbg.pth https://huggingface.co/schirrmacher/ormbg/resolve/main/models/ormbg.pth
+
+echo "Setup completed successfully!"