Spaces:
Sleeping
Sleeping
| import os | |
| import sys | |
| import pathlib | |
| import traceback | |
| import urllib.request | |
| import gradio as gr | |
| import onnxruntime as ort | |
| from huggingface_hub import snapshot_download | |
| # ========================= | |
| # CONFIG — use secrets only | |
| # ========================= | |
| PRIVATE_SPACE_ID = os.getenv("PRIVATE_SPACE_ID") | |
| PRIVATE_PKG_IMPORT = os.getenv("PRIVATE_PKG_IMPORT") | |
| PRIVATE_MODEL_REL = os.getenv("PRIVATE_MODEL_REL") | |
| HF_TOKEN = os.getenv("HF_TOKEN") | |
| if not (PRIVATE_SPACE_ID and PRIVATE_PKG_IMPORT and PRIVATE_MODEL_REL and HF_TOKEN): | |
| raise RuntimeError("Missing one or more required secrets") | |
| # Public fallback | |
| PUBLIC_MODEL_PATH = pathlib.Path(os.getenv("PUBLIC_MODEL_PATH")) | |
| PUBLIC_MODEL_URL = os.getenv("PUBLIC_MODEL_URL") | |
| if not (PUBLIC_MODEL_PATH and PUBLIC_MODEL_URL): | |
| raise RuntimeError("Missing PUBLIC_MODEL_PATH or PUBLIC_MODEL_URL in secrets.") | |
| # ========================= | |
| # DOWNLOAD PRIVATE SNAPSHOT | |
| # ========================= | |
| CACHE_DIR = "private_code" | |
| local_dir = snapshot_download( | |
| repo_id=PRIVATE_SPACE_ID, | |
| repo_type="space", | |
| token=HF_TOKEN, | |
| local_dir=CACHE_DIR, | |
| local_dir_use_symlinks=False, | |
| ) | |
| sys.path.insert(0, local_dir) | |
| # Import private function | |
| try: | |
| mod = __import__(PRIVATE_PKG_IMPORT, fromlist=["remove_background_pil"]) | |
| remove_background_pil = getattr(mod, "remove_background_pil") | |
| except Exception: | |
| raise RuntimeError("Failed to import private function. Ensure PRIVATE_PKG_IMPORT is correct.") | |
| # ========================= | |
| # RESOLVE MODEL PATH | |
| # ========================= | |
| def resolve_model_path() -> str: | |
| private_model = pathlib.Path(local_dir) / PRIVATE_MODEL_REL | |
| if private_model.exists() and private_model.stat().st_size > 5_000_000: | |
| return str(private_model) | |
| PUBLIC_MODEL_PATH.parent.mkdir(parents=True, exist_ok=True) | |
| if not PUBLIC_MODEL_PATH.exists() or PUBLIC_MODEL_PATH.stat().st_size <= 5_000_000: | |
| urllib.request.urlretrieve(PUBLIC_MODEL_URL, PUBLIC_MODEL_PATH) | |
| return str(PUBLIC_MODEL_PATH) | |
| MODEL_PATH = resolve_model_path() | |
| # ========================= | |
| # HEALTH CHECK | |
| # ========================= | |
| def health_check() -> str: | |
| try: | |
| exists = os.path.exists(MODEL_PATH) | |
| size = os.path.getsize(MODEL_PATH) if exists else 0 | |
| providers = ort.get_available_providers() | |
| return f"Model OK: {exists}, size={size}, providers={providers}" | |
| except Exception as e: | |
| return f"Health check failed: {e}" | |
| # ========================= | |
| # INFERENCE | |
| # ========================= | |
| def infer(image, refine=True): | |
| if image is None: | |
| return None | |
| try: | |
| return remove_background_pil(image, model_path=MODEL_PATH, refine=refine) | |
| except Exception: | |
| raise gr.Error("Background removal failed. Please retry.") | |
| # ========================= | |
| # UI | |
| # ========================= | |
| with gr.Blocks() as demo: | |
| gr.Markdown("# 🪄 AI Background Remover (Public Demo)") | |
| with gr.Row(): | |
| with gr.Column(): | |
| inp = gr.Image(type="pil", label="Upload image") | |
| refine = gr.Checkbox(value=True, label="Refine edges") | |
| run_btn = gr.Button("Remove Background", variant="primary") | |
| diag_btn = gr.Button("Run Health Check") | |
| with gr.Column(): | |
| out = gr.Image(type="pil", label="Result (PNG with alpha)") | |
| diag = gr.Code(label="Diagnostics", interactive=False) | |
| run_btn.click(fn=infer, inputs=[inp, refine], outputs=out) | |
| diag_btn.click(fn=health_check, inputs=None, outputs=diag) | |
| if __name__ == "__main__": | |
| demo.launch() |