Update app.py

app.py

@@ -43,7 +43,6 @@ def ensure_rgba(pil: Image.Image) -> Image.Image:
 
 
 def make_checkerboard(w: int, h: int, block: int = 16) -> Image.Image:
-    # Neutral checkerboard
     cols = int(math.ceil(w / block))
     rows = int(math.ceil(h / block))
     board = np.zeros((rows * block, cols * block, 3), dtype=np.uint8)
@@ -79,6 +78,11 @@ def now_ms() -> float:
     return time.perf_counter() * 1000.0
 
 
+def get_device() -> str:
+    """Get device at runtime (important for ZeroGPU)."""
+    return "cuda" if torch.cuda.is_available() else "cpu"
+
+
 @dataclass
 class Timing:
     preprocess_ms: float
@@ -109,7 +113,7 @@ class ModelManager:
       5) IS-Net (isnet-general-use) via rembg
     """
     def __init__(self):
-        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        # NOTE: Don't cache device here - ZeroGPU allocates GPU later
         self._inspy: Optional[Remover] = None
 
         self._torch_models: Dict[str, AutoModelForImageSegmentation] = {}
@@ -132,17 +136,16 @@ class ModelManager:
             pass
 
     def _maybe_sync(self):
-        if self.device == "cuda":
+        if get_device() == "cuda":
             torch.cuda.synchronize()
 
     def _load_inspy(self) -> Remover:
         if self._inspy is None:
-            # jit=False like your sample
             self._inspy = Remover(jit=False)
         return self._inspy
 
     def _offload_torch_models_from_gpu(self, keep_name: str):
-        if self.device != "cuda":
+        if get_device() != "cuda":
             return
         if self._torch_model_on_gpu and self._torch_model_on_gpu != keep_name:
             prev = self._torch_models.get(self._torch_model_on_gpu)
@@ -167,7 +170,7 @@ class ModelManager:
 
         m = AutoModelForImageSegmentation.from_pretrained(model_id, trust_remote_code=True)
         m.eval()
-        # Keep on CPU initially; move to GPU on-demand to avoid T4 OOM.
+        # Keep on CPU initially; move to GPU on-demand
         m.to("cpu")
         self._torch_models[key] = m
         return m
@@ -179,27 +182,28 @@ class ModelManager:
         if name in self._rembg_sessions:
             return self._rembg_sessions[name]
 
-        # Prefer CUDA provider if onnxruntime-gpu is installed; otherwise CPU works.
-        # rembg will pass this into onnxruntime internally.
-        providers = None
+        # Prefer CUDA provider if available
+        providers = ["CUDAExecutionProvider", "CPUExecutionProvider"]
+        
         try:
-            providers = ["CUDAExecutionProvider", "CPUExecutionProvider"]
+            sess = new_session(name, providers=providers)
         except Exception:
-            providers = None
-
-        sess = new_session(name, providers=providers) if providers else new_session(name)
+            # 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 (alpha matte-ish).
+        Runs a torch segmentation model that returns a single-channel mask.
         Returns RGBA (with alpha).
         """
+        device = get_device()  # Check device at runtime!
         m = self._load_torch_model(model_key)
 
         # Put model on GPU for inference if possible
-        if self.device == "cuda":
+        if device == "cuda":
             self._offload_torch_models_from_gpu(keep_name=model_key)
             if self._torch_model_on_gpu != model_key:
                 m.to("cuda")
@@ -209,10 +213,10 @@ class ModelManager:
         orig_size = image_rgb.size
 
         x = self._tf_1024(image_rgb).unsqueeze(0)
-        x = x.to(self.device)
+        x = x.to(device)
 
         with torch.inference_mode():
-            if self.device == "cuda":
+            if device == "cuda":
                 with torch.autocast(device_type="cuda", dtype=torch.float16):
                     preds = m(x)[-1].sigmoid()
             else:
@@ -220,7 +224,7 @@ class ModelManager:
 
         # Convert prediction to PIL alpha channel
         pred = preds[0].squeeze().detach().float().cpu()
-        alpha = transforms.ToPILImage()(pred).resize(orig_size)
+        alpha = transforms.ToPILImage()(pred).resize(orig_size, Image.BILINEAR)
 
         out = image_rgb.convert("RGBA")
         out.putalpha(alpha)
@@ -244,7 +248,6 @@ class ModelManager:
         inf0 = now_ms()
         if model_name == "InSPyReNet":
             remover = self._load_inspy()
-            # The library returns various modes; we want alpha mask and apply ourselves for consistent output
             mask = remover.process(input_image, type="map")
             if isinstance(mask, Image.Image):
                 mask = mask.convert("L")
@@ -259,17 +262,18 @@ class ModelManager:
 
         elif model_name == "U2Net":
             sess = self._get_rembg_session("u2net")
-            # rembg returns bytes (PNG RGBA)
-            out_bytes = rembg_remove(img_rgb, session=sess)
-            out = Image.open(io.BytesIO(out_bytes)).convert("RGBA")
+            # FIX: rembg returns PIL Image when given PIL Image, not bytes!
+            out = rembg_remove(img_rgb, session=sess)
+            out = ensure_rgba(out)
 
         elif model_name == "BRIA RMBG 2.0":
             out = self._run_torch_alpha_model("bria_rmbg_2", img_rgb)
 
         elif model_name == "IS-Net":
             sess = self._get_rembg_session("isnet-general-use")
-            out_bytes = rembg_remove(img_rgb, session=sess)
-            out = Image.open(io.BytesIO(out_bytes)).convert("RGBA")
+            # 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}")
@@ -314,27 +318,20 @@ def run_single(model_name: str, image: Image.Image):
     if image is None:
         return None, None, "Upload an image first.", None
 
-    # Warmup-ish for fairer timing (tiny; avoids huge overhead in UI)
-    # Note: real benchmark tab does proper warmups.
     out_rgba, timing = MANAGER.run(model_name, image)
 
-    # Slider wants (processed, original) or (after, before) depending on component;
-    # we’ll show: left=original, right=on-checkerboard preview of transparent output.
     preview = rgba_on_checkerboard(out_rgba)
-
     out_path = save_temp_png(out_rgba)
     return (image, preview), out_rgba, timing.to_text(), out_path
 
 
 def list_bench_images() -> List[str]:
-    # Put your 10–15 images under bench/
     exts = ("*.jpg", "*.jpeg", "*.png", "*.webp")
     files = []
     for e in exts:
         files += glob.glob(os.path.join("bench", e))
     files = sorted(files)
 
-    # Fallback to repo-root examples like your sample Space
     if not files:
         fallback = []
         for f in ["1.jpg", "2.jpg", "3.png", "4.webp"]:
@@ -348,7 +345,8 @@ def list_bench_images() -> List[str]:
 def run_benchmark(model_name: str, repeats: int = 1):
     files = list_bench_images()
     if not files:
-        return gr.Dataframe(value=[]), "No benchmark images found. Add 10–15 images under bench/."
+        # 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")
@@ -363,12 +361,12 @@ def run_benchmark(model_name: str, repeats: int = 1):
         img = Image.open(f).convert("RGB")
         for r in range(repeats):
             out, timing = MANAGER.run(model_name, img)
-            rows.append({
-                "file": os.path.basename(f),
-                "repeat": r + 1,
-                "total_ms": round(timing.total_ms, 2),
-                "inference_ms": round(timing.inference_ms, 2),
-            })
+            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
 
@@ -380,26 +378,21 @@ def run_benchmark(model_name: str, repeats: int = 1):
         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: {'T4 GPU' if torch.cuda.is_available() else 'CPU'}"
+        f"Device: {'GPU' if torch.cuda.is_available() else 'CPU'}"
     )
 
-    df = gr.Dataframe(
-        headers=["file", "repeat", "total_ms", "inference_ms"],
-        value=[[r["file"], r["repeat"], r["total_ms"], r["inference_ms"]] for r in rows],
-        datatype=["str", "number", "number", "number"],
-        interactive=False
-    )
-    return df, summary
+    # FIX: Return the data directly, not a gr.Dataframe component
+    return rows, summary
 
 
 # ----------------------------
 # UI
 # ----------------------------
 
-with gr.Blocks(title="Background Removal Benchmark (T4)") as demo:
+with gr.Blocks(title="Background Removal Benchmark") as demo:
     gr.Markdown(
         """
-# Background Removal Benchmark (T4)
+# Background Removal Benchmark
 
 Benchmarked models:
 1) InSPyReNet  
@@ -437,7 +430,7 @@ Benchmarked models:
         with gr.Row():
             with gr.Column(scale=1):
                 bench_model = gr.Dropdown(choices=MODEL_CHOICES, value="InSPyReNet", label="Model")
-                repeats = gr.Slider(1, 5, value=1, step=1, label="Repeats per image (higher = more stable averages)")
+                repeats = gr.Slider(1, 5, value=1, step=1, label="Repeats per image")
                 bench_btn = gr.Button("Run benchmark", variant="primary")
             with gr.Column(scale=2):
                 bench_table = gr.Dataframe(
@@ -453,7 +446,6 @@ Benchmarked models:
             outputs=[bench_table, bench_summary]
         )
 
-    # Examples (optional) — if these files exist, they show up like your sample Space
     example_files = []
     for f in ["1.jpg", "2.jpg", "3.png", "4.webp"]:
         if os.path.exists(f):
@@ -466,4 +458,4 @@ Benchmarked models:
         )
 
 if __name__ == "__main__":
-    demo.launch(show_error=True)
+    demo.launch(show_error=True)