Update app.py

app.py

@@ -1,5 +1,4 @@
 import os
-import io
 import time
 import glob
 import math
@@ -15,11 +14,12 @@ from PIL import Image
 import torch
 from torchvision import transforms
 from transformers import AutoModelForImageSegmentation
+from huggingface_hub import hf_hub_download
 
-# Slider component (same one BRIA Space uses)
+# Slider component
 from gradio_imageslider import ImageSlider
 
-# InSPyReNet wrapper (same approach as your sample Space)
+# InSPyReNet wrapper
 from transparent_background import Remover
 
 # rembg (U2Net + IS-Net via ONNX)
@@ -46,17 +46,12 @@ def make_checkerboard(w: int, h: int, block: int = 16) -> Image.Image:
     cols = int(math.ceil(w / block))
     rows = int(math.ceil(h / block))
     board = np.zeros((rows * block, cols * block, 3), dtype=np.uint8)
-
-    c1 = np.array([235, 235, 235], dtype=np.uint8)
-    c2 = np.array([200, 200, 200], dtype=np.uint8)
-
+    c1, c2 = np.array([235, 235, 235], dtype=np.uint8), np.array([200, 200, 200], dtype=np.uint8)
     for r in range(rows):
         for c in range(cols):
             color = c1 if (r + c) % 2 == 0 else c2
             board[r * block:(r + 1) * block, c * block:(c + 1) * block] = color
-
-    board = board[:h, :w, :]
-    return Image.fromarray(board, mode="RGB")
+    return Image.fromarray(board[:h, :w, :], mode="RGB")
 
 
 def rgba_on_checkerboard(rgba: Image.Image) -> Image.Image:
@@ -104,32 +99,19 @@ class Timing:
 # ----------------------------
 
 class ModelManager:
-    """
-    Loads and runs:
-      1) InSPyReNet via transparent_background.Remover()
-      2) BiRefNet via AutoModelForImageSegmentation("ZhengPeng7/BiRefNet", trust_remote_code=True)
-      3) U2Net via rembg (onnxruntime; can use CUDA provider if available)
-      4) BRIA RMBG 2.0 via AutoModelForImageSegmentation("briaai/RMBG-2.0", trust_remote_code=True)
-      5) IS-Net (isnet-general-use) via rembg
-    """
     def __init__(self):
-        # NOTE: Don't cache device here - ZeroGPU allocates GPU later
         self._inspy: Optional[Remover] = None
-
-        self._torch_models: Dict[str, AutoModelForImageSegmentation] = {}
+        self._torch_models: Dict[str, torch.nn.Module] = {}
         self._torch_model_on_gpu: Optional[str] = None
-
-        # rembg sessions
         self._rembg_sessions: Dict[str, object] = {}
+        self._model_load_errors: Dict[str, str] = {}
 
-        # Common transforms for BiRefNet / BRIA RMBG inference
         self._tf_1024 = transforms.Compose([
             transforms.Resize((1024, 1024)),
             transforms.ToTensor(),
             transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
         ])
 
-        # Try to set matmul precision nicely
         try:
             torch.set_float32_matmul_precision("high")
         except Exception:
@@ -154,55 +136,69 @@ class ModelManager:
             self._torch_model_on_gpu = None
             torch.cuda.empty_cache()
 
-    def _load_torch_model(self, key: str) -> AutoModelForImageSegmentation:
-        """
-        key in {"birefnet", "bria_rmbg_2"}
-        """
+    def _load_torch_model(self, key: str) -> torch.nn.Module:
+        """Load BiRefNet or BRIA RMBG 2.0 model."""
         if key in self._torch_models:
             return self._torch_models[key]
+        
+        if key in self._model_load_errors:
+            raise RuntimeError(self._model_load_errors[key])
 
-        if key == "birefnet":
-            model_id = "ZhengPeng7/BiRefNet"
-        elif key == "bria_rmbg_2":
-            model_id = "briaai/RMBG-2.0"
-        else:
-            raise ValueError(f"Unknown torch model key: {key}")
+        model_configs = {
+            "birefnet": "ZhengPeng7/BiRefNet",
+            "bria_rmbg_2": "briaai/RMBG-2.0",
+        }
+        
+        if key not in model_configs:
+            raise ValueError(f"Unknown model key: {key}")
 
-        m = AutoModelForImageSegmentation.from_pretrained(model_id, trust_remote_code=True)
-        m.eval()
-        # Keep on CPU initially; move to GPU on-demand
-        m.to("cpu")
-        self._torch_models[key] = m
-        return m
+        model_id = model_configs[key]
+        
+        try:
+            m = AutoModelForImageSegmentation.from_pretrained(
+                model_id, 
+                trust_remote_code=True
+            )
+            m.eval()
+            m.to("cpu")
+            self._torch_models[key] = m
+            return m
+        except OSError as e:
+            error_msg = str(e)
+            if "gated" in error_msg.lower() or "401" in error_msg or "access" in error_msg.lower():
+                self._model_load_errors[key] = (
+                    f"Model '{model_id}' requires license acceptance.\n"
+                    f"1. Go to https://huggingface.co/{model_id}\n"
+                    f"2. Accept the license agreement\n"
+                    f"3. Add HF_TOKEN secret to your Space settings"
+                )
+            else:
+                self._model_load_errors[key] = f"Failed to load {model_id}: {error_msg}"
+            raise RuntimeError(self._model_load_errors[key])
+        except ImportError as e:
+            self._model_load_errors[key] = (
+                f"Import error loading {model_id}: {e}\n"
+                f"Make sure 'timm' is in requirements.txt"
+            )
+            raise RuntimeError(self._model_load_errors[key])
 
     def _get_rembg_session(self, name: str):
-        """
-        name: "u2net" or "isnet-general-use"
-        """
         if name in self._rembg_sessions:
             return self._rembg_sessions[name]
 
-        # Prefer CUDA provider if available
         providers = ["CUDAExecutionProvider", "CPUExecutionProvider"]
-        
         try:
             sess = new_session(name, providers=providers)
         except Exception:
-            # Fallback to default providers
             sess = new_session(name)
         
         self._rembg_sessions[name] = sess
         return sess
 
     def _run_torch_alpha_model(self, model_key: str, image_rgb: Image.Image) -> Image.Image:
-        """
-        Runs a torch segmentation model that returns a single-channel mask.
-        Returns RGBA (with alpha).
-        """
-        device = get_device()  # Check device at runtime!
+        device = get_device()
         m = self._load_torch_model(model_key)
 
-        # Put model on GPU for inference if possible
         if device == "cuda":
             self._offload_torch_models_from_gpu(keep_name=model_key)
             if self._torch_model_on_gpu != model_key:
@@ -212,8 +208,7 @@ class ModelManager:
         image_rgb = pil_to_rgb(image_rgb)
         orig_size = image_rgb.size
 
-        x = self._tf_1024(image_rgb).unsqueeze(0)
-        x = x.to(device)
+        x = self._tf_1024(image_rgb).unsqueeze(0).to(device)
 
         with torch.inference_mode():
             if device == "cuda":
@@ -222,7 +217,6 @@ class ModelManager:
             else:
                 preds = m(x)[-1].sigmoid()
 
-        # Convert prediction to PIL alpha channel
         pred = preds[0].squeeze().detach().float().cpu()
         alpha = transforms.ToPILImage()(pred).resize(orig_size, Image.BILINEAR)
 
@@ -231,21 +225,19 @@ class ModelManager:
         return out
 
     def run(self, model_name: str, input_image: Image.Image) -> Tuple[Image.Image, Timing]:
-        """
-        Returns (output_rgba, timing).
-        """
         if input_image is None:
             raise ValueError("No input image")
 
         t0 = now_ms()
 
-        # --- preprocess ---
+        # Preprocess
         pre0 = now_ms()
         img_rgb = pil_to_rgb(input_image)
         pre1 = now_ms()
 
-        # --- inference ---
+        # Inference
         inf0 = now_ms()
+        
         if model_name == "InSPyReNet":
             remover = self._load_inspy()
             mask = remover.process(input_image, type="map")
@@ -253,7 +245,6 @@ class ModelManager:
                 mask = mask.convert("L")
             else:
                 mask = Image.fromarray((mask * 255).astype(np.uint8), mode="L")
-
             out = img_rgb.convert("RGBA")
             out.putalpha(mask)
 
@@ -262,7 +253,6 @@ class ModelManager:
 
         elif model_name == "U2Net":
             sess = self._get_rembg_session("u2net")
-            # FIX: rembg returns PIL Image when given PIL Image, not bytes!
             out = rembg_remove(img_rgb, session=sess)
             out = ensure_rgba(out)
 
@@ -271,18 +261,16 @@ class ModelManager:
 
         elif model_name == "IS-Net":
             sess = self._get_rembg_session("isnet-general-use")
-            # FIX: rembg returns PIL Image when given PIL Image, not bytes!
             out = rembg_remove(img_rgb, session=sess)
             out = ensure_rgba(out)
 
         else:
             raise ValueError(f"Unknown model: {model_name}")
 
-        # Make sure GPU timing is accurate
         self._maybe_sync()
         inf1 = now_ms()
 
-        # --- postprocess ---
+        # Postprocess
         post0 = now_ms()
         out = ensure_rgba(out)
         post1 = now_ms()
@@ -318,11 +306,15 @@ def run_single(model_name: str, image: Image.Image):
     if image is None:
         return None, None, "Upload an image first.", None
 
-    out_rgba, timing = MANAGER.run(model_name, image)
-
-    preview = rgba_on_checkerboard(out_rgba)
-    out_path = save_temp_png(out_rgba)
-    return (image, preview), out_rgba, timing.to_text(), out_path
+    try:
+        out_rgba, timing = MANAGER.run(model_name, image)
+        preview = rgba_on_checkerboard(out_rgba)
+        out_path = save_temp_png(out_rgba)
+        return (image, preview), out_rgba, timing.to_text(), out_path
+    except RuntimeError as e:
+        return None, None, f"Error: {str(e)}", None
+    except Exception as e:
+        return None, None, f"Unexpected error: {str(e)}", None
 
 
 def list_bench_images() -> List[str]:
@@ -333,11 +325,9 @@ def list_bench_images() -> List[str]:
     files = sorted(files)
 
     if not files:
-        fallback = []
         for f in ["1.jpg", "2.jpg", "3.png", "4.webp"]:
             if os.path.exists(f):
-                fallback.append(f)
-        files = fallback
+                files.append(f)
     return files
 
 
@@ -345,44 +335,47 @@ def list_bench_images() -> List[str]:
 def run_benchmark(model_name: str, repeats: int = 1):
     files = list_bench_images()
     if not files:
-        # FIX: Return data values, not gr.Dataframe component
         return [], "No benchmark images found. Add 10–15 images under bench/."
 
-    # Warmup: 2 runs on first image (not timed)
-    warm_img = Image.open(files[0]).convert("RGB")
-    for _ in range(2):
-        _ = MANAGER.run(model_name, warm_img)
-
-    rows = []
-    total_ms = 0.0
-    n_images = 0
-
-    for f in files:
-        img = Image.open(f).convert("RGB")
-        for r in range(repeats):
-            out, timing = MANAGER.run(model_name, img)
-            rows.append([
-                os.path.basename(f),
-                r + 1,
-                round(timing.total_ms, 2),
-                round(timing.inference_ms, 2),
-            ])
-            total_ms += timing.total_ms
-            n_images += 1
-
-    avg_ms = total_ms / max(1, n_images)
-    ips = 1000.0 / avg_ms if avg_ms > 0 else 0.0
-
-    summary = (
-        f"Model: {model_name}\n"
-        f"Images: {len(files)} (repeats={repeats}) => runs={n_images}\n"
-        f"Avg total: {avg_ms:.2f} ms\n"
-        f"Estimated throughput: {ips:.2f} images/sec\n"
-        f"Device: {'GPU' if torch.cuda.is_available() else 'CPU'}"
-    )
-
-    # FIX: Return the data directly, not a gr.Dataframe component
-    return rows, summary
+    try:
+        # Warmup
+        warm_img = Image.open(files[0]).convert("RGB")
+        for _ in range(2):
+            _ = MANAGER.run(model_name, warm_img)
+
+        rows = []
+        total_ms = 0.0
+        n_images = 0
+
+        for f in files:
+            img = Image.open(f).convert("RGB")
+            for r in range(repeats):
+                out, timing = MANAGER.run(model_name, img)
+                rows.append([
+                    os.path.basename(f),
+                    r + 1,
+                    round(timing.total_ms, 2),
+                    round(timing.inference_ms, 2),
+                ])
+                total_ms += timing.total_ms
+                n_images += 1
+
+        avg_ms = total_ms / max(1, n_images)
+        ips = 1000.0 / avg_ms if avg_ms > 0 else 0.0
+
+        summary = (
+            f"Model: {model_name}\n"
+            f"Images: {len(files)} (repeats={repeats}) => runs={n_images}\n"
+            f"Avg total: {avg_ms:.2f} ms\n"
+            f"Estimated throughput: {ips:.2f} images/sec\n"
+            f"Device: {'GPU' if torch.cuda.is_available() else 'CPU'}"
+        )
+        return rows, summary
+    
+    except RuntimeError as e:
+        return [], f"Error: {str(e)}"
+    except Exception as e:
+        return [], f"Unexpected error: {str(e)}"
 
 
 # ----------------------------
@@ -395,16 +388,18 @@ with gr.Blocks(title="Background Removal Benchmark") as demo:
 # Background Removal Benchmark
 
 Benchmarked models:
-1) InSPyReNet  
-2) BiRefNet  
-3) U2Net  
-4) BRIA RMBG 2.0  
-5) IS-Net (isnet-general-use)
+1. **InSPyReNet** — transparent-background library  
+2. **BiRefNet** — ZhengPeng7/BiRefNet (requires `timm`)  
+3. **U2Net** — via rembg/ONNX  
+4. **BRIA RMBG 2.0** — briaai/RMBG-2.0 (requires license acceptance)  
+5. **IS-Net** — isnet-general-use via rembg
 
 **Notes**
-- Output download is a true transparent PNG (RGBA).
-- The slider preview composites the transparent result over a checkerboard for visibility.
-- For the benchmark tab, add **10–15 images** under `bench/` in your Space repo.
+- Output is true transparent PNG (RGBA)
+- Slider preview shows result on checkerboard
+- For benchmarks, add images under `bench/` folder
+
+⚠️ **BRIA RMBG 2.0**: Requires accepting license at [huggingface.co/briaai/RMBG-2.0](https://huggingface.co/briaai/RMBG-2.0) and adding `HF_TOKEN` secret to Space settings.
 """
     )
 
@@ -417,7 +412,7 @@ Benchmarked models:
             with gr.Column(scale=2):
                 slider = ImageSlider(label="Before / After", type="pil")
                 out_img = gr.Image(type="pil", label="Output (RGBA)", height=420)
-                timing_box = gr.Textbox(label="Timing", lines=5)
+                timing_box = gr.Textbox(label="Timing / Errors", lines=5)
                 out_file = gr.File(label="Download PNG (transparent)")
 
         run_btn.click(
@@ -451,11 +446,7 @@ Benchmarked models:
         if os.path.exists(f):
             example_files.append([f, "InSPyReNet"])
     if example_files:
-        gr.Examples(
-            examples=example_files,
-            inputs=[inp, model],
-            label="Examples"
-        )
+        gr.Examples(examples=example_files, inputs=[inp, model], label="Examples")
 
 if __name__ == "__main__":
     demo.launch(show_error=True)