Update app.py

app.py

@@ -1,188 +1,109 @@
 import os
 import glob
+import site
 import gradio as gr
 import spaces          # required for ZeroGPU
 import torch           # noqa: F401  (also required by `spaces` for its CUDA patching)
 import onnxruntime as ort
 import numpy as np
-from PIL import Image, ImageDraw
+from PIL import Image
 import cv2
 import tempfile
 
+MODEL_PATH = "birefnet-general.onnx"
+
+
 # -----------------------------------------------------------------
-# FIX: onnxruntime's CUDAExecutionProvider needs libcublasLt.so /
-# libcudnn.so etc. at runtime. Those files ARE installed (they ship
-# inside the `nvidia-cublas-cu12` / `nvidia-cudnn-cu12` packages that
-# come with torch) but they live under site-packages/nvidia/*/lib,
-# which is NOT on the system's shared-library search path — so the
-# CUDA provider fails to load with "libcublasLt.so.12: cannot open
-# shared object file" and onnxruntime silently falls back to CPU.
-# We point the dynamic linker at the right folders ourselves, and
-# also ask onnxruntime to preload what it can find. This must run
-# BEFORE any InferenceSession with CUDAExecutionProvider is created.
+# Infra fix (keep this): points the dynamic linker at the CUDA libs
+# so onnxruntime can actually use the GPU instead of silently
+# falling back to CPU. Nothing to do with image editing — leave as is.
 # -----------------------------------------------------------------
 def _fix_cuda_library_path():
+    lib_dirs = set()
     try:
         import nvidia
-        nvidia_root = os.path.dirname(nvidia.__file__)
-        lib_dirs = glob.glob(os.path.join(nvidia_root, "*", "lib"))
-        if lib_dirs:
-            existing = os.environ.get("LD_LIBRARY_PATH", "")
-            os.environ["LD_LIBRARY_PATH"] = ":".join(lib_dirs) + (":" + existing if existing else "")
-            print("[ORT] Added to LD_LIBRARY_PATH:", lib_dirs)
+        for root in list(nvidia.__path__):
+            lib_dirs.update(glob.glob(os.path.join(root, "*", "lib")))
     except Exception as e:
-        print("[ORT] Could not locate nvidia pip lib dirs:", e)
+        print("[ORT] nvidia namespace import lookup failed:", e)
+
+    if not lib_dirs:
+        search_roots = []
+        try:
+            search_roots += site.getsitepackages()
+        except Exception:
+            pass
+        try:
+            search_roots.append(site.getusersitepackages())
+        except Exception:
+            pass
+        for root in search_roots:
+            lib_dirs.update(glob.glob(os.path.join(root, "nvidia", "*", "lib")))
+
+    if lib_dirs:
+        lib_dirs = sorted(lib_dirs)
+        existing = os.environ.get("LD_LIBRARY_PATH", "")
+        os.environ["LD_LIBRARY_PATH"] = ":".join(lib_dirs) + (":" + existing if existing else "")
+        print("[ORT] Added to LD_LIBRARY_PATH:", lib_dirs)
+    else:
+        print("[ORT] WARNING: no nvidia CUDA lib directories found on disk.")
 
     try:
-        ort.preload_dlls()  # onnxruntime >= 1.21; harmless no-op-ish if it fails
-        print("[ORT] preload_dlls() succeeded")
+        ort.preload_dlls()
+        print("[ORT] preload_dlls() called")
     except Exception as e:
         print("[ORT] preload_dlls() unavailable/failed (non-fatal):", e)
 
 
 _fix_cuda_library_path()
 
-MODEL_PATH = "birefnet-general.onnx"
-
-# -----------------------------------------------------------------
-# IMPORTANT: do NOT create the CUDA InferenceSession here at import
-# time. ZeroGPU only attaches a real GPU to the process while a
-# @spaces.GPU-decorated function is executing — outside of that, this
-# process has no GPU, so building a CUDAExecutionProvider session here
-# would fail. Instead we build it lazily, from inside the decorated
-# function below. Note: ZeroGPU may run each call in a fresh worker,
-# so this "cache" mainly exists for the case where it doesn't — don't
-# rely on it for performance, just correctness.
-# -----------------------------------------------------------------
+# Lazy session — built the first time it's needed, inside the GPU call
+# (ZeroGPU only attaches a real GPU while @spaces.GPU is running).
 _session = None
-_input_name = None
-_output_name = None
 
 
 def get_session():
-    global _session, _input_name, _output_name
+    global _session
     if _session is None:
         _session = ort.InferenceSession(
             MODEL_PATH,
             providers=["CUDAExecutionProvider", "CPUExecutionProvider"]
         )
-        _input_name = _session.get_inputs()[0].name
-        _output_name = _session.get_outputs()[0].name
         print("[ORT] active providers:", _session.get_providers())
-    return _session, _input_name, _output_name
-
+    return _session
 
 
-@spaces.GPU(duration=30)  # bump this if a cold-start (first call / model load) needs more time
+@spaces.GPU(duration=30)
 def run_inference(inp):
-    session, input_name, output_name = get_session()
+    session = get_session()
+    input_name = session.get_inputs()[0].name
+    output_name = session.get_outputs()[0].name
     return session.run([output_name], {input_name: inp})[0]
 
 
-def remove_bg(
-    image,
-    show_box=False,
-    canvas_size=1200,
-    padding=15
-):
+def remove_bg(image):
     img = np.array(image.convert("RGB"))
     h, w = img.shape[:2]
-    inp = cv2.resize(img, (1024, 1024))
-    inp = inp.astype(np.float32) / 255.0
+
+    inp = cv2.resize(img, (1024, 1024)).astype(np.float32) / 255.0
     inp = np.transpose(inp, (2, 0, 1))
     inp = np.expand_dims(inp, 0)
 
     pred = run_inference(inp)
+    mask = cv2.resize(pred[0][0], (w, h))
+    alpha = (np.clip(mask, 0, 1) * 255).astype(np.uint8)
 
-    mask = pred[0][0]
-    mask = cv2.resize(mask, (w, h))
-    mask = np.clip(mask, 0, 1)
-    alpha = (mask * 255).astype(np.uint8)
-    ys, xs = np.where(alpha > 20)
-    if len(xs) == 0 or len(ys) == 0:
-        raise Exception("No object detected")
-    x1 = int(xs.min())
-    x2 = int(xs.max())
-    y1 = int(ys.min())
-    y2 = int(ys.max())
-    # -----------------------------------
-    # Crop object
-    # -----------------------------------
-    obj_rgb = img[y1:y2 + 1, x1:x2 + 1]
-    obj_alpha = alpha[y1:y2 + 1, x1:x2 + 1]
-    obj_rgba = np.dstack([obj_rgb, obj_alpha])
-    obj = Image.fromarray(obj_rgba)
-    obj_w = obj.width
-    obj_h = obj.height
-    # -----------------------------------
-    # Scale while preserving MIN padding
-    # -----------------------------------
-
-    max_w = canvas_size - (padding * 2)
-    max_h = canvas_size - (padding * 2)
-
-    scale_w = max_w / obj_w
-    scale_h = max_h / obj_h
-
-    scale = min(scale_w, scale_h)
-
-    new_w = round(obj_w * scale)
-    new_h = round(obj_h * scale)
-
-    obj = obj.resize(
-        (new_w, new_h),
-        Image.LANCZOS
-    )
-    # -----------------------------------
-    # White canvas
-    # -----------------------------------
-    canvas = Image.new(
-        "RGBA",
-        (canvas_size, canvas_size),
-        (255, 255, 255, 255)
-    )
-    paste_x = (canvas_size - new_w) // 2
-    paste_y = (canvas_size - new_h) // 2
-    canvas.paste(
-        obj,
-        (paste_x, paste_y),
-        obj
-    )
-    result = canvas.convert("RGB")
-    # -----------------------------------
-    # Optional bounding box
-    # -----------------------------------
-    if show_box:
-        draw = ImageDraw.Draw(result)
-        draw.rectangle(
-            [
-                paste_x,
-                paste_y,
-                paste_x + new_w,
-                paste_y + new_h
-            ],
-            outline=(0, 100, 255),
-            width=4
-        )
-    tmp = tempfile.NamedTemporaryFile(
-        suffix=".png",
-        delete=False
-    )
-    result.save(
-        tmp.name,
-        format="PNG"
-    )
+    rgba = np.dstack([img, alpha])
+    result = Image.fromarray(rgba)
+
+    tmp = tempfile.NamedTemporaryFile(suffix=".png", delete=False)
+    result.save(tmp.name, format="PNG")
     return tmp.name
 
 
 demo = gr.Interface(
     fn=remove_bg,
-    inputs=[
-        gr.Image(type="pil"),
-        gr.Checkbox(value=False),
-        gr.Number(value=1200),
-        gr.Number(value=15)
-    ],
+    inputs=gr.Image(type="pil"),
     outputs=gr.File(),
     api_name="remove_bg"
 )