feat: Add BEN2 background removal functionality and integrate into SnapForm

- Introduced BEN2 model for background removal in the panel and properties.
- Updated runtime to handle background removal commands and model downloads.
- Implemented background removal logic in worker, including image processing and saving results.
- Enhanced tests to cover new background removal features and ensure proper integration.
- Updated version to 0.9.28 to reflect new features and improvements.

src/snapform/worker.py

@@ -188,6 +188,7 @@ def download_model(args):
         repo_id=args.repository,
         local_dir=model_dir,
         tqdm_class=progress_class,
+        allow_patterns=args.allow_pattern or None,
     )
     write_status(args.status, 1.0, "Model is ready")
 
@@ -202,6 +203,76 @@ def select_device(requested, torch):
     return torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 
+def infer_foreground(input_path, model_dir, device, torch):
+    from ben2 import BEN_Base
+    from PIL import Image
+
+    input_path = Path(input_path)
+    if input_path.suffix.lower() not in IMAGE_SUFFIXES:
+        raise RuntimeError("BEN2 requires a JPG, JPEG, PNG, or WEBP image")
+    load_started = time.perf_counter()
+    model = BEN_Base.from_pretrained(model_dir)
+    model = model.to(device=device)
+    model.eval()
+    load_ms = (time.perf_counter() - load_started) * 1000.0
+
+    with Image.open(input_path) as opened_image:
+        image = opened_image.convert("RGB")
+    if device.type == "cuda":
+        torch.cuda.synchronize()
+    inference_started = time.perf_counter()
+    with torch.inference_mode():
+        foreground = model.inference(
+            image,
+            refine_foreground=False,
+        )
+    if device.type == "cuda":
+        torch.cuda.synchronize()
+    inference_ms = (time.perf_counter() - inference_started) * 1000.0
+    del model
+    if device.type == "cuda":
+        torch.cuda.empty_cache()
+    return foreground, load_ms, inference_ms
+
+
+def remove_background(args):
+    total_started = time.perf_counter()
+    write_status(args.status, 0.02, "Starting BEN2 Runtime")
+
+    import torch
+
+    output_directory = Path(args.output_dir)
+    output_directory.mkdir(parents=True, exist_ok=False)
+    device = select_device(args.device, torch)
+
+    write_status(args.status, 0.12, f"Loading BEN2 on {device.type}")
+    write_status(args.status, 0.35, "Removing background")
+    foreground, load_ms, inference_ms = infer_foreground(
+        args.input,
+        args.model_dir,
+        device,
+        torch,
+    )
+
+    write_status(args.status, 0.92, "Writing transparent PNG")
+    output_path = output_directory / "foreground.png"
+    foreground.convert("RGBA").save(
+        output_path,
+        format="PNG",
+        compress_level=1,
+    )
+    if not output_path.is_file():
+        raise RuntimeError("BEN2 did not write a foreground image")
+
+    total_ms = (time.perf_counter() - total_started) * 1000.0
+    message = (
+        f"Background removed in {inference_ms:.0f} ms "
+        f"(model load {load_ms:.0f} ms, total {total_ms:.0f} ms)"
+    )
+    print(message)
+    write_status(args.status, 1.0, message)
+
+
 def exported_paths(output_directory, export_format):
     output_directory = Path(output_directory)
     if export_format == "glb":
@@ -246,6 +317,30 @@ def reconstruct(args):
 
     try:
         device = select_device(args.device, torch)
+        remove_background_model_dir = getattr(
+            args,
+            "remove_background_model_dir",
+            "",
+        )
+        if remove_background_model_dir:
+            write_status(args.status, 0.12, "Removing image background")
+            foreground, load_ms, inference_ms = infer_foreground(
+                images[0],
+                remove_background_model_dir,
+                device,
+                torch,
+            )
+            foreground.convert("RGBA").save(
+                output_directory / "foreground.png",
+                format="PNG",
+                compress_level=1,
+            )
+            if not (output_directory / "foreground.png").is_file():
+                raise RuntimeError("BEN2 did not write a foreground image")
+            print(
+                f"BEN2 inference {inference_ms:.0f} ms; "
+                f"model load {load_ms:.0f} ms"
+            )
         write_status(args.status, 0.15, f"Loading model on {device.type}")
         model = DepthAnything3.from_pretrained(args.model_dir)
         model = model.to(device=device)
@@ -311,6 +406,18 @@ def parser():
     download.add_argument("--repository", required=True)
     download.add_argument("--model-dir", required=True)
     download.add_argument("--status", required=True)
+    download.add_argument("--allow-pattern", action="append", default=[])
+
+    background = commands.add_parser("remove-background")
+    background.add_argument("--input", required=True)
+    background.add_argument("--model-dir", required=True)
+    background.add_argument("--output-dir", required=True)
+    background.add_argument("--status", required=True)
+    background.add_argument(
+        "--device",
+        choices=("auto", "cuda", "cpu"),
+        required=True,
+    )
 
     reconstruction = commands.add_parser("reconstruct")
     reconstruction.add_argument("--input", required=True)
@@ -322,6 +429,7 @@ def parser():
         required=True,
     )
     reconstruction.add_argument("--status", required=True)
+    reconstruction.add_argument("--remove-background-model-dir", default="")
     reconstruction.add_argument("--process-res", type=int, required=True)
     reconstruction.add_argument(
         "--process-res-method",
@@ -370,8 +478,10 @@ def main():
     try:
         if args.command == "download":
             download_model(args)
-        else:
+        elif args.command == "reconstruct":
             reconstruct(args)
+        else:
+            remove_background(args)
     except Exception as error:
         write_status(args.status, 1.0, "Task failed", str(error))
         raise