Update app.py

Returning the two received masks along with the generated images for debugging purposes

app.py

@@ -1,285 +1,3 @@
-# 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":
@@ -304,6 +22,10 @@ def run_local(base_image, base_mask, reference_image, ref_mask, seed, base_mask_
     ref_image = ref_image.convert("RGB")
     ref_mask = ref_mask.convert("L")
 
+    # Store the received masks for return
+    received_tar_mask = tar_mask.copy()
+    received_ref_mask = ref_mask.copy()
+
     return_ref_mask = ref_mask.copy()
 
     tar_image = np.asarray(tar_image)
@@ -390,81 +112,6 @@ def run_local(base_image, base_mask, reference_image, ref_mask, seed, base_mask_
     edited_image = Image.fromarray(edited_image)
 
     if ref_mask_option != "Label to Mask":
-        return [show_diptych_ref_tar, edited_image]
-    else:
-        return [return_ref_mask, show_diptych_ref_tar, edited_image]
-
-
-def update_ui(option):
-    if option == "Draw Mask":
-        return gr.update(visible=False), gr.update(visible=True)
+        return [show_diptych_ref_tar, edited_image, received_tar_mask, received_ref_mask]
     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()
+        return [return_ref_mask, show_diptych_ref_tar, edited_image, received_tar_mask, received_ref_mask]