isat commited on
Commit
ecf3916
·
verified ·
1 Parent(s): 00f7e20

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +358 -1
app.py CHANGED
@@ -1,3 +1,285 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
  @spaces.GPU
2
  def run_local(base_image, base_mask, reference_image, ref_mask, seed, base_mask_option, ref_mask_option, text_prompt):
3
  if base_mask_option == "Draw Mask":
@@ -114,4 +396,79 @@ def run_local(base_image, base_mask, reference_image, ref_mask, seed, base_mask_
114
  if ref_mask_option != "Label to Mask":
115
  return [show_diptych_ref_tar, edited_image, received_tar_mask, received_ref_mask]
116
  else:
117
- return [return_ref_mask, show_diptych_ref_tar, edited_image, received_tar_mask, received_ref_mask]
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # app.py — storage-safe + HF Hub friendly + SAM import guard
2
+
3
+ import os
4
+
5
+ # ---------- ENV & THREADS (set BEFORE importing numpy/torch) ----------
6
+ omp_val = (
7
+ os.getenv("OMP_NUM_THREADS")
8
+ or os.getenv("OMP-NUM-THREADS")
9
+ or os.getenv("OMPNUMTHREADS")
10
+ or "2"
11
+ )
12
+ try:
13
+ omp_val = str(int(omp_val))
14
+ except Exception:
15
+ omp_val = "2"
16
+ os.environ["OMP_NUM_THREADS"] = omp_val # must be a positive integer string
17
+
18
+ # Persistent caches
19
+ os.environ.setdefault("HF_HOME", "/data/.huggingface")
20
+ os.environ.setdefault("HF_HUB_CACHE", "/data/.huggingface/hub")
21
+ os.environ.setdefault("HF_DATASETS_CACHE", "/data/.huggingface/datasets")
22
+ # (TRANSFORMERS_CACHE is deprecated; rely on HF_HOME) # https://huggingface.co/docs/huggingface_hub/en/guides/manage-cache
23
+
24
+ # Disable Xet path, enable fast transfer
25
+ os.environ.setdefault("HF_HUB_DISABLE_XET", "1")
26
+ os.environ.setdefault("HF_HUB_ENABLE_HF_TRANSFER", "1")
27
+
28
+ # ---------- NOW safe to import heavy libs ----------
29
+ import sys
30
+ import cv2
31
+ import numpy as np
32
+ import torch
33
+ import gradio as gr
34
+ from PIL import Image, ImageFilter, ImageDraw
35
+
36
+ try:
37
+ torch.set_num_threads(int(omp_val))
38
+ torch.set_num_interop_threads(1)
39
+ except Exception:
40
+ pass
41
+
42
+ # ---------- HUB IMPORTS ----------
43
+ from huggingface_hub import snapshot_download, hf_hub_download
44
+ from diffusers import FluxFillPipeline, FluxPriorReduxPipeline
45
+
46
+ import math
47
+ from utils.utils import (
48
+ get_bbox_from_mask, expand_bbox, pad_to_square, box2squre, crop_back, expand_image_mask
49
+ )
50
+
51
+ # ---------- Ensure GroundingDINO & SAM are the right ones ----------
52
+ def _ensure_local_editable(pkg_name, rel_path):
53
+ try:
54
+ __import__(pkg_name)
55
+ except ImportError:
56
+ os.system(f"{sys.executable} -m pip install -e {rel_path}")
57
+
58
+ # GroundingDINO (local editable if present)
59
+ _ensure_local_editable("GroundingDINO", "GroundingDINO")
60
+
61
+ # SAM: verify the real package; fix automatically if a wrong one is installed
62
+ def _ensure_official_sam():
63
+ try:
64
+ import segment_anything as sa
65
+ if not hasattr(sa, "sam_model_registry"):
66
+ raise ImportError("Found 'segment_anything' without sam_model_registry")
67
+ except Exception:
68
+ # Nuke imposters and install the official repo
69
+ os.system(f"{sys.executable} -m pip uninstall -y segment-anything segment_anything")
70
+ os.system(f"{sys.executable} -m pip install -U git+https://github.com/facebookresearch/segment-anything.git")
71
+
72
+ _ensure_official_sam()
73
+
74
+ # Now import
75
+ sys.path.append(os.path.join(os.getcwd(), "GroundingDINO"))
76
+ import torchvision
77
+ from GroundingDINO.groundingdino.util.inference import load_model
78
+ from segment_anything import sam_model_registry, SamPredictor # official API
79
+ import spaces
80
+ import GroundingDINO.groundingdino.datasets.transforms as T
81
+ from GroundingDINO.groundingdino.util.utils import clean_state_dict, get_phrases_from_posmap
82
+
83
+ # ---------- PATHS ----------
84
+ PERSIST_ROOT = "/data"
85
+ MODELS_DIR = os.path.join(PERSIST_ROOT, "models")
86
+ CKPT_DIR = os.path.join(PERSIST_ROOT, "checkpoints")
87
+ os.makedirs(MODELS_DIR, exist_ok=True)
88
+ os.makedirs(CKPT_DIR, exist_ok=True)
89
+
90
+ # GroundingDINO config and checkpoint
91
+ GROUNDING_DINO_CONFIG_PATH = "./GroundingDINO_SwinB.cfg.py"
92
+ GROUNDING_DINO_CHECKPOINT_PATH = os.path.join(CKPT_DIR, "groundingdino_swinb_cogcoor.pth")
93
+
94
+ # Segment-Anything checkpoint
95
+ SAM_ENCODER_VERSION = "vit_h"
96
+ SAM_CHECKPOINT_PATH = os.path.join(CKPT_DIR, "sam_vit_h_4b8939.pth")
97
+
98
+ # ---------- AUTH TOKEN ----------
99
+ hf_token = os.getenv("HF_TOKEN") or os.getenv("HUGGINGFACE_HUB_TOKEN")
100
+
101
+ # ---------- DOWNLOAD CHECKPOINTS (single files) ----------
102
+ # 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
103
+ if not os.path.exists(GROUNDING_DINO_CHECKPOINT_PATH):
104
+ g_dino_file = hf_hub_download(
105
+ repo_id="ShilongLiu/GroundingDINO",
106
+ filename="groundingdino_swinb_cogcoor.pth",
107
+ local_dir=CKPT_DIR,
108
+ token=hf_token,
109
+ )
110
+ if g_dino_file != GROUNDING_DINO_CHECKPOINT_PATH:
111
+ os.replace(g_dino_file, GROUNDING_DINO_CHECKPOINT_PATH)
112
+
113
+ if not os.path.exists(SAM_CHECKPOINT_PATH):
114
+ sam_file = hf_hub_download(
115
+ repo_id="mrtlive/segment-anything-model", # remove "spaces/"
116
+ repo_type="space", # tell the Hub it's a Space
117
+ filename="sam_vit_h_4b8939.pth",
118
+ local_dir=CKPT_DIR,
119
+ token=hf_token,
120
+ )
121
+ if sam_file != SAM_CHECKPOINT_PATH:
122
+ os.replace(sam_file, SAM_CHECKPOINT_PATH)
123
+
124
+ # ---------- DOWNLOAD MODELS (filtered snapshots into /data) ----------
125
+ FILL_DIR = os.path.join(MODELS_DIR, "FLUX.1-Fill-dev")
126
+ REDUX_DIR = os.path.join(MODELS_DIR, "FLUX.1-Redux-dev")
127
+ LORA_DIR = os.path.join(MODELS_DIR, "insertanything_model")
128
+ for path in (FILL_DIR, REDUX_DIR, LORA_DIR):
129
+ os.makedirs(path, exist_ok=True)
130
+
131
+ # Only pull what we need (weights/configs). Keep symlinks to avoid copies.
132
+ if not os.listdir(FILL_DIR):
133
+ snapshot_download(
134
+ repo_id="black-forest-labs/FLUX.1-Fill-dev",
135
+ local_dir=FILL_DIR,
136
+ local_dir_use_symlinks=True,
137
+ allow_patterns=["*.safetensors", "*.json", "*.yaml", "*.txt", "*.py", "*.model"],
138
+ token=hf_token,
139
+ )
140
+
141
+ if not os.listdir(REDUX_DIR):
142
+ snapshot_download(
143
+ repo_id="black-forest-labs/FLUX.1-Redux-dev",
144
+ local_dir=REDUX_DIR,
145
+ local_dir_use_symlinks=True,
146
+ allow_patterns=["*.safetensors", "*.json", "*.yaml", "*.txt", "*.py", "*.model"],
147
+ token=hf_token,
148
+ )
149
+
150
+ if not os.listdir(LORA_DIR):
151
+ snapshot_download(
152
+ repo_id="WensongSong/Insert-Anything",
153
+ local_dir=LORA_DIR,
154
+ local_dir_use_symlinks=True,
155
+ allow_patterns=["*.safetensors", "*.json", "*.yaml", "*.txt"],
156
+ token=hf_token,
157
+ )
158
+
159
+ # ---------- BUILD MODELS ----------
160
+ # GroundingDINO
161
+ groundingdino_model = load_model(
162
+ model_config_path=GROUNDING_DINO_CONFIG_PATH,
163
+ model_checkpoint_path=GROUNDING_DINO_CHECKPOINT_PATH,
164
+ device="cuda"
165
+ )
166
+
167
+ # SAM + Predictor (registry API from official SAM) # https://github.com/facebookresearch/segment-anything
168
+ sam = sam_model_registry[SAM_ENCODER_VERSION](checkpoint=SAM_CHECKPOINT_PATH)
169
+ sam.to(device="cuda")
170
+ sam_predictor = SamPredictor(sam)
171
+
172
+ # Diffusers (Flux)
173
+ dtype = torch.bfloat16
174
+ size = (768, 768)
175
+
176
+ pipe = FluxFillPipeline.from_pretrained(
177
+ FILL_DIR,
178
+ torch_dtype=dtype
179
+ ).to("cuda")
180
+
181
+ pipe.load_lora_weights(
182
+ os.path.join(LORA_DIR, "20250321_steps5000_pytorch_lora_weights.safetensors")
183
+ )
184
+
185
+ redux = FluxPriorReduxPipeline.from_pretrained(REDUX_DIR).to(dtype=dtype).to("cuda")
186
+
187
+ # ---------- APP LOGIC ----------
188
+ def transform_image(image_pil):
189
+ transform = T.Compose(
190
+ [
191
+ T.RandomResize([800], max_size=1333),
192
+ T.ToTensor(),
193
+ T.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
194
+ ]
195
+ )
196
+ image, _ = transform(image_pil, None) # 3, h, w
197
+ return image
198
+
199
+
200
+ def get_grounding_output(model, image, caption, box_threshold=0.25, text_threshold=0.25, with_logits=True):
201
+ caption = caption.lower().strip()
202
+ if not caption.endswith("."):
203
+ caption = caption + "."
204
+ with torch.no_grad():
205
+ outputs = model(image[None], captions=[caption])
206
+ logits = outputs["pred_logits"].cpu().sigmoid()[0] # (nq, 256)
207
+ boxes = outputs["pred_boxes"].cpu()[0] # (nq, 4)
208
+
209
+ # filter output
210
+ filt_mask = logits.max(dim=1)[0] > box_threshold
211
+ logits_filt = logits[filt_mask]
212
+ boxes_filt = boxes[filt_mask]
213
+
214
+ # get phrase
215
+ tokenlizer = model.tokenizer
216
+ tokenized = tokenlizer(caption)
217
+ pred_phrases, scores = [], []
218
+ for logit, box in zip(logits_filt, boxes_filt):
219
+ pred_phrase = get_phrases_from_posmap(logit > text_threshold, tokenized, tokenlizer)
220
+ pred_phrases.append(pred_phrase + f"({str(logit.max().item())[:4]})" if with_logits else pred_phrase)
221
+ scores.append(logit.max().item())
222
+ return boxes_filt, torch.Tensor(scores), pred_phrases
223
+
224
+
225
+ def get_mask(image, label):
226
+ global groundingdino_model, sam_predictor
227
+ image_pil = image.convert("RGB")
228
+ transformed_image = transform_image(image_pil)
229
+
230
+ boxes_filt, scores, pred_phrases = get_grounding_output(
231
+ groundingdino_model, transformed_image, label
232
+ )
233
+
234
+ W, H = image_pil.size
235
+ for i in range(boxes_filt.size(0)):
236
+ boxes_filt[i] = boxes_filt[i] * torch.Tensor([W, H, W, H])
237
+ boxes_filt[i][:2] -= boxes_filt[i][2:] / 2
238
+ boxes_filt[i][2:] += boxes_filt[i][:2]
239
+ boxes_filt = boxes_filt.cpu()
240
+
241
+ nms_idx = torchvision.ops.nms(boxes_filt, scores, 0.8).numpy().tolist()
242
+ boxes_filt = boxes_filt[nms_idx]
243
+
244
+ image_np = np.array(image_pil)
245
+ sam_predictor.set_image(image_np)
246
+ transformed_boxes = sam_predictor.transform.apply_boxes_torch(
247
+ boxes_filt, image_np.shape[:2]
248
+ ).to("cuda")
249
+
250
+ masks, _, _ = sam_predictor.predict_torch(
251
+ point_coords=None,
252
+ point_labels=None,
253
+ boxes=transformed_boxes,
254
+ multimask_output=False,
255
+ )
256
+ result_mask = masks[0][0].cpu().numpy()
257
+ return Image.fromarray(result_mask)
258
+
259
+
260
+ def create_highlighted_mask(image_np, mask_np, alpha=0.5, gray_value=128):
261
+ if mask_np.max() <= 1.0:
262
+ mask_np = (mask_np * 255).astype(np.uint8)
263
+ mask_bool = mask_np > 128
264
+ image_float = image_np.astype(np.float32)
265
+ gray_overlay = np.full_like(image_float, gray_value, dtype=np.float32)
266
+ result = image_float.copy()
267
+ result[mask_bool] = (1 - alpha) * image_float[mask_bool] + alpha * gray_overlay[mask_bool]
268
+ return result.astype(np.uint8)
269
+
270
+
271
+ # ---------- EXAMPLES ----------
272
+ ref_dir = './examples/ref_image'
273
+ ref_mask_dir = './examples/ref_mask'
274
+ image_dir = './examples/source_image'
275
+ image_mask_dir = './examples/source_mask'
276
+
277
+ ref_list = sorted([os.path.join(ref_dir, f) for f in os.listdir(ref_dir) if f.lower().endswith((".jpg", ".png", ".jpeg"))])
278
+ 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"))])
279
+ image_list = sorted([os.path.join(image_dir, f) for f in os.listdir(image_dir) if f.lower().endswith((".jpg", ".png", ".jpeg"))])
280
+ 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"))])
281
+
282
+
283
  @spaces.GPU
284
  def run_local(base_image, base_mask, reference_image, ref_mask, seed, base_mask_option, ref_mask_option, text_prompt):
285
  if base_mask_option == "Draw Mask":
 
396
  if ref_mask_option != "Label to Mask":
397
  return [show_diptych_ref_tar, edited_image, received_tar_mask, received_ref_mask]
398
  else:
399
+ return [return_ref_mask, show_diptych_ref_tar, edited_image, received_tar_mask, received_ref_mask]
400
+
401
+
402
+ def update_ui(option):
403
+ if option == "Draw Mask":
404
+ return gr.update(visible=False), gr.update(visible=True)
405
+ else:
406
+ return gr.update(visible=True), gr.update(visible=False)
407
+
408
+
409
+ with gr.Blocks() as demo:
410
+ gr.Markdown("# Insert-Anything")
411
+ gr.Markdown("### Make sure to select the correct mask button!!")
412
+ gr.Markdown("### Click the output image to toggle between Diptych and final results!!")
413
+
414
+ with gr.Row():
415
+ with gr.Column(scale=1):
416
+ with gr.Row():
417
+ base_image = gr.ImageEditor(label="Background Image", sources="upload", type="pil",
418
+ brush=gr.Brush(colors=["#FFFFFF"], default_size=30, color_mode="fixed"),
419
+ layers=False, interactive=True)
420
+ base_mask = gr.ImageEditor(label="Background Mask", sources="upload", type="pil",
421
+ layers=False, brush=False, eraser=False)
422
+ with gr.Row():
423
+ base_mask_option = gr.Radio(["Draw Mask", "Upload with Mask"], label="Background Mask Input Option",
424
+ value="Upload with Mask")
425
+
426
+ with gr.Row():
427
+ ref_image = gr.ImageEditor(label="Reference Image", sources="upload", type="pil",
428
+ brush=gr.Brush(colors=["#FFFFFF"], default_size=30, color_mode="fixed"),
429
+ layers=False, interactive=True)
430
+ ref_mask = gr.ImageEditor(label="Reference Mask", sources="upload", type="pil",
431
+ layers=False, brush=False, eraser=False)
432
+
433
+ with gr.Row():
434
+ ref_mask_option = gr.Radio(["Draw Mask", "Upload with Mask", "Label to Mask"],
435
+ label="Reference Mask Input Option", value="Upload with Mask")
436
+ with gr.Row():
437
+ text_prompt = gr.Textbox(label="Label",
438
+ placeholder="Enter the category of the reference object, e.g., car, dress, toy, etc.")
439
+
440
+ with gr.Column(scale=1):
441
+ baseline_gallery = gr.Gallery(label='Output', show_label=True, elem_id="gallery", height=695, columns=1)
442
+ with gr.Accordion("Advanced Option", open=True):
443
+ seed = gr.Slider(label="Seed", minimum=-1, maximum=999_999_999, step=1, value=666)
444
+ gr.Markdown("### Guidelines")
445
+ gr.Markdown(" Users can try using different seeds. For example, seeds like 42 and 123456 may produce different effects.")
446
+ gr.Markdown(" Draw Mask means manually drawing a mask on the original image.")
447
+ gr.Markdown(" Upload with Mask means uploading a mask file.")
448
+ gr.Markdown(" Label to Mask means simply inputting a label to automatically extract the mask and obtain the result.")
449
+
450
+ run_local_button = gr.Button(value="Run")
451
+
452
+ # examples
453
+ num_examples = len(image_list)
454
+ for i in range(num_examples):
455
+ with gr.Row():
456
+ if i == 0:
457
+ gr.Examples([image_list[i]], inputs=[base_image], label="Examples - Background Image", examples_per_page=1)
458
+ gr.Examples([image_mask_list[i]], inputs=[base_mask], label="Examples - Background Mask", examples_per_page=1)
459
+ gr.Examples([ref_list[i]], inputs=[ref_image], label="Examples - Reference Object", examples_per_page=1)
460
+ gr.Examples([ref_mask_list[i]], inputs=[ref_mask], label="Examples - Reference Mask", examples_per_page=1)
461
+ else:
462
+ gr.Examples([image_list[i]], inputs=[base_image], examples_per_page=1, label="")
463
+ gr.Examples([image_mask_list[i]], inputs=[base_mask], examples_per_page=1, label="")
464
+ gr.Examples([ref_list[i]], inputs=[ref_image], examples_per_page=1, label="")
465
+ gr.Examples([ref_mask_list[i]], inputs=[ref_mask], examples_per_page=1, label="")
466
+ if i < num_examples - 1:
467
+ gr.HTML("<hr>")
468
+
469
+ run_local_button.click(
470
+ fn=run_local,
471
+ inputs=[base_image, base_mask, ref_image, ref_mask, seed, base_mask_option, ref_mask_option, text_prompt],
472
+ outputs=[baseline_gallery]
473
+ )
474
+ demo.launch()