Update app.py

app.py

@@ -1,3 +1,285 @@
+# app.py — storage-safe + HF Hub friendly + SAM import guard
+
+import os
+
+# ---------- ENV & THREADS (set BEFORE importing numpy/torch) ----------
+omp_val = (
+    os.getenv("OMP_NUM_THREADS")
+    or os.getenv("OMP-NUM-THREADS")
+    or os.getenv("OMPNUMTHREADS")
+    or "2"
+)
+try:
+    omp_val = str(int(omp_val))
+except Exception:
+    omp_val = "2"
+os.environ["OMP_NUM_THREADS"] = omp_val  # must be a positive integer string
+
+# 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")
+# (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")
+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
+
+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
+from diffusers import FluxFillPipeline, FluxPriorReduxPipeline
+
+import math
+from utils.utils import (
+    get_bbox_from_mask, expand_bbox, pad_to_square, box2squre, crop_back, expand_image_mask
+)
+
+# ---------- Ensure GroundingDINO & SAM are the right ones ----------
+def _ensure_local_editable(pkg_name, rel_path):
+    try:
+        __import__(pkg_name)
+    except ImportError:
+        os.system(f"{sys.executable} -m pip install -e {rel_path}")
+
+# GroundingDINO (local editable if present)
+_ensure_local_editable("GroundingDINO", "GroundingDINO")
+
+# 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")
+
+_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"
+MODELS_DIR = os.path.join(PERSIST_ROOT, "models")
+CKPT_DIR = os.path.join(PERSIST_ROOT, "checkpoints")
+os.makedirs(MODELS_DIR, exist_ok=True)
+os.makedirs(CKPT_DIR, exist_ok=True)
+
+# GroundingDINO config and checkpoint
+GROUNDING_DINO_CONFIG_PATH = "./GroundingDINO_SwinB.cfg.py"
+GROUNDING_DINO_CHECKPOINT_PATH = os.path.join(CKPT_DIR, "groundingdino_swinb_cogcoor.pth")
+
+# Segment-Anything checkpoint
+SAM_ENCODER_VERSION = "vit_h"
+SAM_CHECKPOINT_PATH = os.path.join(CKPT_DIR, "sam_vit_h_4b8939.pth")
+
+# ---------- AUTH TOKEN ----------
+hf_token = os.getenv("HF_TOKEN") or os.getenv("HUGGINGFACE_HUB_TOKEN")
+
+# ---------- DOWNLOAD CHECKPOINTS (single files) ----------
+# 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",
+        filename="groundingdino_swinb_cogcoor.pth",
+        local_dir=CKPT_DIR,
+        token=hf_token,
+    )
+    if g_dino_file != GROUNDING_DINO_CHECKPOINT_PATH:
+        os.replace(g_dino_file, GROUNDING_DINO_CHECKPOINT_PATH)
+
+if not os.path.exists(SAM_CHECKPOINT_PATH):
+    sam_file = hf_hub_download(
+        repo_id="mrtlive/segment-anything-model",   # remove "spaces/"
+        repo_type="space",                          # tell the Hub it's a Space
+        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)
+
+# ---------- DOWNLOAD MODELS (filtered snapshots into /data) ----------
+FILL_DIR = os.path.join(MODELS_DIR, "FLUX.1-Fill-dev")
+REDUX_DIR = os.path.join(MODELS_DIR, "FLUX.1-Redux-dev")
+LORA_DIR = os.path.join(MODELS_DIR, "insertanything_model")
+for path in (FILL_DIR, REDUX_DIR, LORA_DIR):
+    os.makedirs(path, exist_ok=True)
+
+# Only pull what we need (weights/configs). Keep symlinks to avoid copies.
+if not os.listdir(FILL_DIR):
+    snapshot_download(
+        repo_id="black-forest-labs/FLUX.1-Fill-dev",
+        local_dir=FILL_DIR,
+        local_dir_use_symlinks=True,
+        allow_patterns=["*.safetensors", "*.json", "*.yaml", "*.txt", "*.py", "*.model"],
+        token=hf_token,
+    )
+
+if not os.listdir(REDUX_DIR):
+    snapshot_download(
+        repo_id="black-forest-labs/FLUX.1-Redux-dev",
+        local_dir=REDUX_DIR,
+        local_dir_use_symlinks=True,
+        allow_patterns=["*.safetensors", "*.json", "*.yaml", "*.txt", "*.py", "*.model"],
+        token=hf_token,
+    )
+
+if not os.listdir(LORA_DIR):
+    snapshot_download(
+        repo_id="WensongSong/Insert-Anything",
+        local_dir=LORA_DIR,
+        local_dir_use_symlinks=True,
+        allow_patterns=["*.safetensors", "*.json", "*.yaml", "*.txt"],
+        token=hf_token,
+    )
+
+# ---------- BUILD MODELS ----------
+# GroundingDINO
+groundingdino_model = load_model(
+    model_config_path=GROUNDING_DINO_CONFIG_PATH,
+    model_checkpoint_path=GROUNDING_DINO_CHECKPOINT_PATH,
+    device="cuda"
+)
+
+# 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 (Flux)
+dtype = torch.bfloat16
+size = (768, 768)
+
+pipe = FluxFillPipeline.from_pretrained(
+    FILL_DIR,
+    torch_dtype=dtype
+).to("cuda")
+
+pipe.load_lora_weights(
+    os.path.join(LORA_DIR, "20250321_steps5000_pytorch_lora_weights.safetensors")
+)
+
+redux = FluxPriorReduxPipeline.from_pretrained(REDUX_DIR).to(dtype=dtype).to("cuda")
+
+# ---------- APP LOGIC ----------
+def transform_image(image_pil):
+    transform = T.Compose(
+        [
+            T.RandomResize([800], max_size=1333),
+            T.ToTensor(),
+            T.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
+        ]
+    )
+    image, _ = transform(image_pil, None)  # 3, h, w
+    return image
+
+
+def get_grounding_output(model, image, caption, box_threshold=0.25, text_threshold=0.25, with_logits=True):
+    caption = caption.lower().strip()
+    if not caption.endswith("."):
+        caption = caption + "."
+    with torch.no_grad():
+        outputs = model(image[None], captions=[caption])
+    logits = outputs["pred_logits"].cpu().sigmoid()[0]  # (nq, 256)
+    boxes = outputs["pred_boxes"].cpu()[0]              # (nq, 4)
+
+    # filter output
+    filt_mask = logits.max(dim=1)[0] > box_threshold
+    logits_filt = logits[filt_mask]
+    boxes_filt = boxes[filt_mask]
+
+    # get phrase
+    tokenlizer = model.tokenizer
+    tokenized = tokenlizer(caption)
+    pred_phrases, scores = [], []
+    for logit, box in zip(logits_filt, boxes_filt):
+        pred_phrase = get_phrases_from_posmap(logit > text_threshold, tokenized, tokenlizer)
+        pred_phrases.append(pred_phrase + f"({str(logit.max().item())[:4]})" if with_logits else pred_phrase)
+        scores.append(logit.max().item())
+    return boxes_filt, torch.Tensor(scores), pred_phrases
+
+
+def get_mask(image, label):
+    global groundingdino_model, sam_predictor
+    image_pil = image.convert("RGB")
+    transformed_image = transform_image(image_pil)
+
+    boxes_filt, scores, pred_phrases = get_grounding_output(
+        groundingdino_model, transformed_image, label
+    )
+
+    W, H = image_pil.size
+    for i in range(boxes_filt.size(0)):
+        boxes_filt[i] = boxes_filt[i] * torch.Tensor([W, H, W, H])
+        boxes_filt[i][:2] -= boxes_filt[i][2:] / 2
+        boxes_filt[i][2:] += boxes_filt[i][:2]
+    boxes_filt = boxes_filt.cpu()
+
+    nms_idx = torchvision.ops.nms(boxes_filt, scores, 0.8).numpy().tolist()
+    boxes_filt = boxes_filt[nms_idx]
+
+    image_np = np.array(image_pil)
+    sam_predictor.set_image(image_np)
+    transformed_boxes = sam_predictor.transform.apply_boxes_torch(
+        boxes_filt, image_np.shape[:2]
+    ).to("cuda")
+
+    masks, _, _ = sam_predictor.predict_torch(
+        point_coords=None,
+        point_labels=None,
+        boxes=transformed_boxes,
+        multimask_output=False,
+    )
+    result_mask = masks[0][0].cpu().numpy()
+    return Image.fromarray(result_mask)
+
+
+def create_highlighted_mask(image_np, mask_np, alpha=0.5, gray_value=128):
+    if mask_np.max() <= 1.0:
+        mask_np = (mask_np * 255).astype(np.uint8)
+    mask_bool = mask_np > 128
+    image_float = image_np.astype(np.float32)
+    gray_overlay = np.full_like(image_float, gray_value, dtype=np.float32)
+    result = image_float.copy()
+    result[mask_bool] = (1 - alpha) * image_float[mask_bool] + alpha * gray_overlay[mask_bool]
+    return result.astype(np.uint8)
+
+
+# ---------- EXAMPLES ----------
+ref_dir = './examples/ref_image'
+ref_mask_dir = './examples/ref_mask'
+image_dir = './examples/source_image'
+image_mask_dir = './examples/source_mask'
+
+ref_list = sorted([os.path.join(ref_dir, f) for f in os.listdir(ref_dir) if f.lower().endswith((".jpg", ".png", ".jpeg"))])
+ref_mask_list = sorted([os.path.join(ref_mask_dir, f) for f in os.listdir(ref_mask_dir) if f.lower().endswith((".jpg", ".png", ".jpeg"))])
+image_list = sorted([os.path.join(image_dir, f) for f in os.listdir(image_dir) if f.lower().endswith((".jpg", ".png", ".jpeg"))])
+image_mask_list = sorted([os.path.join(image_mask_dir, f) for f in os.listdir(image_mask_dir) if f.lower().endswith((".jpg", ".png", ".jpeg"))])
+
+
 @spaces.GPU
 def run_local(base_image, base_mask, reference_image, ref_mask, seed, base_mask_option, ref_mask_option, text_prompt):
     if base_mask_option == "Draw Mask":
@@ -114,4 +396,79 @@ def run_local(base_image, base_mask, reference_image, ref_mask, seed, base_mask_
     if ref_mask_option != "Label to Mask":
         return [show_diptych_ref_tar, edited_image, received_tar_mask, received_ref_mask]
     else:
-        return [return_ref_mask, show_diptych_ref_tar, edited_image, received_tar_mask, received_ref_mask]
+        return [return_ref_mask, show_diptych_ref_tar, edited_image, received_tar_mask, received_ref_mask]
+
+
+def update_ui(option):
+    if option == "Draw Mask":
+        return gr.update(visible=False), gr.update(visible=True)
+    else:
+        return gr.update(visible=True), gr.update(visible=False)
+
+
+with gr.Blocks() as demo:
+    gr.Markdown("# Insert-Anything")
+    gr.Markdown("### Make sure to select the correct mask button!!")
+    gr.Markdown("### Click the output image to toggle between Diptych and final results!!")
+
+    with gr.Row():
+        with gr.Column(scale=1):
+            with gr.Row():
+                base_image = gr.ImageEditor(label="Background Image", sources="upload", type="pil",
+                                            brush=gr.Brush(colors=["#FFFFFF"], default_size=30, color_mode="fixed"),
+                                            layers=False, interactive=True)
+                base_mask = gr.ImageEditor(label="Background Mask", sources="upload", type="pil",
+                                           layers=False, brush=False, eraser=False)
+            with gr.Row():
+                base_mask_option = gr.Radio(["Draw Mask", "Upload with Mask"], label="Background Mask Input Option",
+                                            value="Upload with Mask")
+
+            with gr.Row():
+                ref_image = gr.ImageEditor(label="Reference Image", sources="upload", type="pil",
+                                           brush=gr.Brush(colors=["#FFFFFF"], default_size=30, color_mode="fixed"),
+                                           layers=False, interactive=True)
+                ref_mask = gr.ImageEditor(label="Reference Mask", sources="upload", type="pil",
+                                          layers=False, brush=False, eraser=False)
+
+            with gr.Row():
+                ref_mask_option = gr.Radio(["Draw Mask", "Upload with Mask", "Label to Mask"],
+                                           label="Reference Mask Input Option", value="Upload with Mask")
+            with gr.Row():
+                text_prompt = gr.Textbox(label="Label",
+                                         placeholder="Enter the category of the reference object, e.g., car, dress, toy, etc.")
+
+        with gr.Column(scale=1):
+            baseline_gallery = gr.Gallery(label='Output', show_label=True, elem_id="gallery", height=695, columns=1)
+            with gr.Accordion("Advanced Option", open=True):
+                seed = gr.Slider(label="Seed", minimum=-1, maximum=999_999_999, step=1, value=666)
+                gr.Markdown("### Guidelines")
+                gr.Markdown(" Users can try using different seeds. For example, seeds like 42 and 123456 may produce different effects.")
+                gr.Markdown(" Draw Mask means manually drawing a mask on the original image.")
+                gr.Markdown(" Upload with Mask means uploading a mask file.")
+                gr.Markdown(" Label to Mask means simply inputting a label to automatically extract the mask and obtain the result.")
+
+    run_local_button = gr.Button(value="Run")
+
+    # examples
+    num_examples = len(image_list)
+    for i in range(num_examples):
+        with gr.Row():
+            if i == 0:
+                gr.Examples([image_list[i]], inputs=[base_image], label="Examples - Background Image", examples_per_page=1)
+                gr.Examples([image_mask_list[i]], inputs=[base_mask], label="Examples - Background Mask", examples_per_page=1)
+                gr.Examples([ref_list[i]], inputs=[ref_image], label="Examples - Reference Object", examples_per_page=1)
+                gr.Examples([ref_mask_list[i]], inputs=[ref_mask], label="Examples - Reference Mask", examples_per_page=1)
+            else:
+                gr.Examples([image_list[i]], inputs=[base_image], examples_per_page=1, label="")
+                gr.Examples([image_mask_list[i]], inputs=[base_mask], examples_per_page=1, label="")
+                gr.Examples([ref_list[i]], inputs=[ref_image], examples_per_page=1, label="")
+                gr.Examples([ref_mask_list[i]], inputs=[ref_mask], examples_per_page=1, label="")
+        if i < num_examples - 1:
+            gr.HTML("<hr>")
+
+    run_local_button.click(
+        fn=run_local,
+        inputs=[base_image, base_mask, ref_image, ref_mask, seed, base_mask_option, ref_mask_option, text_prompt],
+        outputs=[baseline_gallery]
+    )
+demo.launch()