Update app.py

app.py

@@ -1,9 +1,8 @@
-# app.py — storage-safe + HF Hub friendly
+# app.py — storage-safe + HF Hub friendly + SAM import guard
 
 import os
 
 # ---------- ENV & THREADS (set BEFORE importing numpy/torch) ----------
-# Accept any of these names from Space Settings, prefer the standard one:
 omp_val = (
     os.getenv("OMP_NUM_THREADS")
     or os.getenv("OMP-NUM-THREADS")
@@ -16,11 +15,11 @@ except Exception:
     omp_val = "2"
 os.environ["OMP_NUM_THREADS"] = omp_val  # must be a positive integer string
 
-# Send all caches to persistent storage
+# Persistent caches
 os.environ.setdefault("HF_HOME", "/data/.huggingface")
 os.environ.setdefault("HF_HUB_CACHE", "/data/.huggingface/hub")
 os.environ.setdefault("HF_DATASETS_CACHE", "/data/.huggingface/datasets")
-# NOTE: TRANSFORMERS_CACHE is deprecated; using HF_HOME instead.
+# (TRANSFORMERS_CACHE is deprecated; rely on HF_HOME)  # https://huggingface.co/docs/huggingface_hub/en/guides/manage-cache
 
 # Disable Xet path, enable fast transfer
 os.environ.setdefault("HF_HUB_DISABLE_XET", "1")
@@ -34,7 +33,6 @@ import torch
 import gradio as gr
 from PIL import Image, ImageFilter, ImageDraw
 
-# Optional: align PyTorch thread pools with OMP setting
 try:
     torch.set_num_threads(int(omp_val))
     torch.set_num_interop_threads(1)
@@ -42,34 +40,45 @@ except Exception:
     pass
 
 # ---------- HUB IMPORTS ----------
-from huggingface_hub import snapshot_download, hf_hub_download  # noqa: E402
-from diffusers import FluxFillPipeline, FluxPriorReduxPipeline  # noqa: E402
+from huggingface_hub import snapshot_download, hf_hub_download
+from diffusers import FluxFillPipeline, FluxPriorReduxPipeline
 
-import math  # noqa: E402
-from utils.utils import (  # noqa: E402
+import math
+from utils.utils import (
     get_bbox_from_mask, expand_bbox, pad_to_square, box2squre, crop_back, expand_image_mask
 )
 
-# Optional editable installs ONLY if import fails (prefer requirements.txt)
+# ---------- Ensure GroundingDINO & SAM are the right ones ----------
 def _ensure_local_editable(pkg_name, rel_path):
     try:
         __import__(pkg_name)
     except ImportError:
-        os.system(f"python -m pip install -e {rel_path}")
+        os.system(f"{sys.executable} -m pip install -e {rel_path}")
 
-_ensure_local_editable("segment_anything", "segment_anything")
+# GroundingDINO (local editable if present)
 _ensure_local_editable("GroundingDINO", "GroundingDINO")
 
-sys.path.append(os.path.join(os.getcwd(), "GroundingDINO"))
-sys.path.append(os.path.join(os.getcwd(), "segment_anything"))
+# SAM: verify the real package; fix automatically if a wrong one is installed
+def _ensure_official_sam():
+    try:
+        import segment_anything as sa
+        if not hasattr(sa, "sam_model_registry"):
+            raise ImportError("Found 'segment_anything' without sam_model_registry")
+    except Exception:
+        # Nuke imposters and install the official repo
+        os.system(f"{sys.executable} -m pip uninstall -y segment-anything segment_anything")
+        os.system(f"{sys.executable} -m pip install -U git+https://github.com/facebookresearch/segment-anything.git")
 
-import torchvision  # noqa: E402
-from GroundingDINO.groundingdino.util.inference import load_model  # noqa: E402
-# Use the stable SAM API (avoids build_sam import error)
-from segment_anything import sam_model_registry, SamPredictor  # noqa: E402
-import spaces  # noqa: E402
-import GroundingDINO.groundingdino.datasets.transforms as T  # noqa: E402
-from GroundingDINO.groundingdino.util.utils import clean_state_dict, get_phrases_from_posmap  # noqa: E402
+_ensure_official_sam()
+
+# Now import
+sys.path.append(os.path.join(os.getcwd(), "GroundingDINO"))
+import torchvision
+from GroundingDINO.groundingdino.util.inference import load_model
+from segment_anything import sam_model_registry, SamPredictor  # official API
+import spaces
+import GroundingDINO.groundingdino.datasets.transforms as T
+from GroundingDINO.groundingdino.util.utils import clean_state_dict, get_phrases_from_posmap
 
 # ---------- PATHS ----------
 PERSIST_ROOT = "/data"
@@ -90,7 +99,7 @@ SAM_CHECKPOINT_PATH = os.path.join(CKPT_DIR, "sam_vit_h_4b8939.pth")
 hf_token = os.getenv("HF_TOKEN") or os.getenv("HUGGINGFACE_HUB_TOKEN")
 
 # ---------- DOWNLOAD CHECKPOINTS (single files) ----------
-# GroundingDINO ckpt (single file)
+# Use hf_hub_download for single files, which returns a cached path. Keep files under /data.  # https://huggingface.co/docs/huggingface_hub/en/guides/download
 if not os.path.exists(GROUNDING_DINO_CHECKPOINT_PATH):
     g_dino_file = hf_hub_download(
         repo_id="ShilongLiu/GroundingDINO",
@@ -101,7 +110,6 @@ if not os.path.exists(GROUNDING_DINO_CHECKPOINT_PATH):
     if g_dino_file != GROUNDING_DINO_CHECKPOINT_PATH:
         os.replace(g_dino_file, GROUNDING_DINO_CHECKPOINT_PATH)
 
-# SAM ckpt (single file)
 if not os.path.exists(SAM_CHECKPOINT_PATH):
     sam_file = hf_hub_download(
         repo_id="spaces/mrtlive/segment-anything-model",
@@ -155,12 +163,12 @@ groundingdino_model = load_model(
     device="cuda"
 )
 
-# SAM + Predictor (registry API)
+# SAM + Predictor (registry API from official SAM)  # https://github.com/facebookresearch/segment-anything
 sam = sam_model_registry[SAM_ENCODER_VERSION](checkpoint=SAM_CHECKPOINT_PATH)
 sam.to(device="cuda")
 sam_predictor = SamPredictor(sam)
 
-# Diffusers
+# Diffusers (Flux)
 dtype = torch.bfloat16
 size = (768, 768)