feat: add BiRefNet HR-matting model download and manifest entry

docker/download_models.py

@@ -186,6 +186,7 @@ REMBG_MODELS = [
     "birefnet-portrait",
     "birefnet-general",
     "birefnet-matting",
+    "birefnet-hr-matting",
 ]
 
 # PaddleOCR PP-OCRv5 HuggingFace model repos to pre-download.
@@ -235,12 +236,59 @@ def _register_birefnet_matting():
     sessions_class.append(BiRefNetMattingSession)
 
 
+def _register_birefnet_hr_matting():
+    """Register BiRefNet HR-matting ONNX session for 2048x2048 high-res matting."""
+    import os
+    import numpy as np
+    import pooch
+    from PIL import Image
+    from rembg.sessions import sessions_class
+    from rembg.sessions.birefnet_general import BiRefNetSessionGeneral
+
+    class BiRefNetHRMattingSession(BiRefNetSessionGeneral):
+        @classmethod
+        def download_models(cls, *args, **kwargs):
+            fname = f"{cls.name(*args, **kwargs)}.onnx"
+            pooch.retrieve(
+                "https://github.com/ZhengPeng7/BiRefNet/releases/download/v1/BiRefNet_HR-matting-epoch_135.onnx",
+                None,
+                fname=fname,
+                path=cls.u2net_home(*args, **kwargs),
+                progressbar=True,
+            )
+            return os.path.join(cls.u2net_home(*args, **kwargs), fname)
+
+        @classmethod
+        def name(cls, *args, **kwargs):
+            return "birefnet-hr-matting"
+
+        def predict(self, img, *args, **kwargs):
+            ort_outs = self.inner_session.run(
+                None,
+                self.normalize(
+                    img, (0.485, 0.456, 0.406), (0.229, 0.224, 0.225), (2048, 2048)
+                ),
+            )
+            pred = ort_outs[0][:, 0, :, :]
+            ma = np.max(pred)
+            mi = np.min(pred)
+            denom = ma - mi
+            pred = (pred - mi) / denom if denom > 0 else pred * 0
+            pred = np.squeeze(pred)
+            mask = Image.fromarray((pred * 255).astype("uint8"), mode="L")
+            mask = mask.resize(img.size, Image.LANCZOS)
+            return [mask]
+
+    sessions_class.append(BiRefNetHRMattingSession)
+
+
 def download_rembg_models():
     """Download all rembg ONNX models."""
     print("=== Downloading rembg models ===")
     from rembg import new_session
 
     _register_birefnet_matting()
+    _register_birefnet_hr_matting()
 
     for model in REMBG_MODELS:
         print(f"  Downloading {model}...")

docker/feature-manifest.json

@@ -7,7 +7,7 @@
     "background-removal": {
       "name": "Background Removal",
       "description": "Remove image backgrounds with AI",
-      "estimatedSize": "3-4 GB",
+      "estimatedSize": "4-5 GB",
       "packages": {
         "common": ["rembg==2.0.62"],
         "amd64": ["onnxruntime-gpu==1.20.1", "mediapipe==0.10.21"],
@@ -42,9 +42,14 @@
           "id": "rembg-birefnet-matting",
           "downloadFn": "rembg_session",
           "args": ["birefnet-matting"]
+        },
+        {
+          "id": "rembg-birefnet-hr-matting",
+          "downloadFn": "rembg_session",
+          "args": ["birefnet-hr-matting"]
         }
       ],
-      "enablesTools": ["remove-background", "passport-photo"]
+      "enablesTools": ["remove-background", "passport-photo", "transparency-fixer"]
     },
     "face-detection": {
       "name": "Face Detection",