Update app.py

app.py

@@ -1,33 +1,46 @@
 # app.py — storage-safe + HF Hub friendly
 
 import os
-import sys
-import cv2
-import numpy as np
-import torch
-import gradio as gr
-from PIL import Image, ImageFilter, ImageDraw
 
-# ---------- ENV & THREADS ----------
-# Map a Spaces variable (no underscores allowed) to the real OpenMP var.
-omp_val = os.getenv("OMP-NUM-THREADS") or os.getenv("OMPNUMTHREADS") or "2"
-os.environ["OMP_NUM_THREADS"] = omp_val
+# ---------- 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")
+    or os.getenv("OMPNUMTHREADS")
+    or "2"
+)
 try:
-    torch.set_num_threads(int(omp_val))
-    torch.set_num_interop_threads(1)
+    omp_val = str(int(omp_val))
 except Exception:
-    pass
+    omp_val = "2"
+os.environ["OMP_NUM_THREADS"] = omp_val  # must be a positive integer string
 
 # Send all caches to persistent storage
 os.environ.setdefault("HF_HOME", "/data/.huggingface")
 os.environ.setdefault("HF_HUB_CACHE", "/data/.huggingface/hub")
-os.environ.setdefault("TRANSFORMERS_CACHE", "/data/.huggingface/transformers")
 os.environ.setdefault("HF_DATASETS_CACHE", "/data/.huggingface/datasets")
+# NOTE: TRANSFORMERS_CACHE is deprecated; using HF_HOME instead.
 
 # Disable Xet path, enable fast transfer
 os.environ.setdefault("HF_HUB_DISABLE_XET", "1")
 os.environ.setdefault("HF_HUB_ENABLE_HF_TRANSFER", "1")
 
+# ---------- NOW safe to import heavy libs ----------
+import sys
+import cv2
+import numpy as np
+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)
+except Exception:
+    pass
+
 # ---------- HUB IMPORTS ----------
 from huggingface_hub import snapshot_download, hf_hub_download  # noqa: E402
 from diffusers import FluxFillPipeline, FluxPriorReduxPipeline  # noqa: E402
@@ -37,7 +50,7 @@ from utils.utils import (  # noqa: E402
     get_bbox_from_mask, expand_bbox, pad_to_square, box2squre, crop_back, expand_image_mask
 )
 
-# Optional editable installs ONLY if import fails (use requirements.txt ideally)
+# Optional editable installs ONLY if import fails (prefer requirements.txt)
 def _ensure_local_editable(pkg_name, rel_path):
     try:
         __import__(pkg_name)
@@ -52,7 +65,8 @@ sys.path.append(os.path.join(os.getcwd(), "segment_anything"))
 
 import torchvision  # noqa: E402
 from GroundingDINO.groundingdino.util.inference import load_model  # noqa: E402
-from segment_anything import build_sam, SamPredictor  # 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
@@ -78,26 +92,25 @@ hf_token = os.getenv("HF_TOKEN") or os.getenv("HUGGINGFACE_HUB_TOKEN")
 # ---------- DOWNLOAD CHECKPOINTS (single files) ----------
 # GroundingDINO ckpt (single file)
 if not os.path.exists(GROUNDING_DINO_CHECKPOINT_PATH):
-    G_DINO_FILE = hf_hub_download(
+    g_dino_file = hf_hub_download(
         repo_id="ShilongLiu/GroundingDINO",
         filename="groundingdino_swinb_cogcoor.pth",
         local_dir=CKPT_DIR,
         token=hf_token,
     )
-    if G_DINO_FILE != GROUNDING_DINO_CHECKPOINT_PATH:
-        # Ensure the expected path exists for later code
-        os.replace(G_DINO_FILE, 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(
+    sam_file = hf_hub_download(
         repo_id="spaces/mrtlive/segment-anything-model",
         filename="sam_vit_h_4b8939.pth",
         local_dir=CKPT_DIR,
         token=hf_token,
     )
-    if SAM_FILE != SAM_CHECKPOINT_PATH:
-        os.replace(SAM_FILE, SAM_CHECKPOINT_PATH)
+    if sam_file != SAM_CHECKPOINT_PATH:
+        os.replace(sam_file, SAM_CHECKPOINT_PATH)
 
 # ---------- DOWNLOAD MODELS (filtered snapshots into /data) ----------
 FILL_DIR = os.path.join(MODELS_DIR, "FLUX.1-Fill-dev")
@@ -142,8 +155,8 @@ groundingdino_model = load_model(
     device="cuda"
 )
 
-# SAM + Predictor
-sam = build_sam(checkpoint=SAM_CHECKPOINT_PATH)
+# SAM + Predictor (registry API)
+sam = sam_model_registry[SAM_ENCODER_VERSION](checkpoint=SAM_CHECKPOINT_PATH)
 sam.to(device="cuda")
 sam_predictor = SamPredictor(sam)