Upload 28 files

.gitattributes

@@ -33,3 +33,15 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+examples/ref_image/1.png filter=lfs diff=lfs merge=lfs -text
+examples/ref_image/2.png filter=lfs diff=lfs merge=lfs -text
+examples/ref_image/3.png filter=lfs diff=lfs merge=lfs -text
+examples/ref_image/4.png filter=lfs diff=lfs merge=lfs -text
+examples/result/1.png filter=lfs diff=lfs merge=lfs -text
+examples/result/2.png filter=lfs diff=lfs merge=lfs -text
+examples/result/3.png filter=lfs diff=lfs merge=lfs -text
+examples/result/4.png filter=lfs diff=lfs merge=lfs -text
+examples/result/5.png filter=lfs diff=lfs merge=lfs -text
+examples/source_image/1.png filter=lfs diff=lfs merge=lfs -text
+examples/source_image/2.png filter=lfs diff=lfs merge=lfs -text
+examples/source_image/3.png filter=lfs diff=lfs merge=lfs -text

app.py

@@ -0,0 +1,245 @@
+import os
+import sys
+import cv2
+import numpy as np
+import torch
+import gradio as gr
+from PIL import Image, ImageFilter, ImageDraw
+from huggingface_hub import snapshot_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
+
+
+hf_token = os.getenv("HF_TOKEN")
+
+snapshot_download(repo_id="black-forest-labs/FLUX.1-Fill-dev", local_dir="./FLUX.1-Fill-dev", token=hf_token)
+snapshot_download(repo_id="black-forest-labs/FLUX.1-Redux-dev", local_dir="./FLUX.1-Redux-dev", token=hf_token)
+snapshot_download(repo_id="WensongSong/Insert-Anything", local_dir="./insertanything_model", token=hf_token)
+
+
+dtype = torch.bfloat16
+size = (768, 768)
+
+pipe = FluxFillPipeline.from_pretrained(
+    "./FLUX.1-Fill-dev",
+    torch_dtype=dtype
+).to("cuda")
+
+pipe.load_lora_weights(
+    "./insertanything_model/20250321-082022_steps5000_pytorch_lora_weights.safetensors"
+)
+
+
+redux = FluxPriorReduxPipeline.from_pretrained("./FLUX.1-Redux-dev").to(dtype=dtype).to("cuda")
+
+
+
+###   example  #####
+ref_dir='./examples/ref_image'
+ref_mask_dir='./examples/ref_mask'
+image_dir='./examples/source_image'
+image_mask_dir='./examples/source_mask'
+
+ref_list=[os.path.join(ref_dir,file) for file in os.listdir(ref_dir) if '.jpg' in file or '.png' in file or '.jpeg' in file ]
+ref_list.sort()
+
+ref_mask_list=[os.path.join(ref_mask_dir,file) for file in os.listdir(ref_mask_dir) if '.jpg' in file or '.png' in file or '.jpeg' in file]
+ref_mask_list.sort()
+
+image_list=[os.path.join(image_dir,file) for file in os.listdir(image_dir) if '.jpg' in file or '.png' in file or '.jpeg' in file ]
+image_list.sort()
+
+image_mask_list=[os.path.join(image_mask_dir,file) for file in os.listdir(image_mask_dir) if '.jpg' in file or '.png' in file or '.jpeg' in file]
+image_mask_list.sort()
+###   example  #####
+
+
+
+
+def run_local(base_image, base_mask, reference_image, ref_mask, seed, base_mask_option, ref_mask_option):
+
+    if base_mask_option == "Draw Mask":
+        tar_image = base_image["image"]
+        tar_mask = base_image["mask"]
+    else:
+        tar_image = base_image["image"]
+        tar_mask = base_mask
+
+    if ref_mask_option == "Draw Mask":
+        ref_image = reference_image["image"]
+        ref_mask = reference_image["mask"]
+    else:
+        ref_image = reference_image["image"]
+        ref_mask = ref_mask
+
+
+    tar_image = tar_image.convert("RGB")
+    tar_mask = tar_mask.convert("L")
+    ref_image = ref_image.convert("RGB")
+    ref_mask = ref_mask.convert("L")
+
+    tar_image = np.asarray(tar_image)
+    tar_mask = np.asarray(tar_mask)
+    tar_mask = np.where(tar_mask > 128, 1, 0).astype(np.uint8)
+
+    ref_image = np.asarray(ref_image)
+    ref_mask = np.asarray(ref_mask)
+    ref_mask = np.where(ref_mask > 128, 1, 0).astype(np.uint8)
+
+
+    ref_box_yyxx = get_bbox_from_mask(ref_mask)
+    ref_mask_3 = np.stack([ref_mask,ref_mask,ref_mask],-1)
+    masked_ref_image = ref_image * ref_mask_3 + np.ones_like(ref_image) * 255 * (1-ref_mask_3) 
+    y1,y2,x1,x2 = ref_box_yyxx
+    masked_ref_image = masked_ref_image[y1:y2,x1:x2,:] 
+    ref_mask = ref_mask[y1:y2,x1:x2] 
+    ratio = 1.3
+    masked_ref_image, ref_mask = expand_image_mask(masked_ref_image, ref_mask, ratio=ratio) 
+
+
+    masked_ref_image = pad_to_square(masked_ref_image, pad_value = 255, random = False) 
+
+    kernel = np.ones((7, 7), np.uint8)
+    iterations = 2
+    tar_mask = cv2.dilate(tar_mask, kernel, iterations=iterations)
+
+    # zome in
+    tar_box_yyxx = get_bbox_from_mask(tar_mask)
+    tar_box_yyxx = expand_bbox(tar_mask, tar_box_yyxx, ratio=1.2)
+
+    tar_box_yyxx_crop =  expand_bbox(tar_image, tar_box_yyxx, ratio=2)    #1.2 1.6
+    tar_box_yyxx_crop = box2squre(tar_image, tar_box_yyxx_crop) # crop box
+    y1,y2,x1,x2 = tar_box_yyxx_crop
+
+
+    old_tar_image = tar_image.copy()
+    tar_image = tar_image[y1:y2,x1:x2,:]
+    tar_mask = tar_mask[y1:y2,x1:x2]
+
+    H1, W1 = tar_image.shape[0], tar_image.shape[1]
+    # zome in
+
+
+    tar_mask = pad_to_square(tar_mask, pad_value=0)
+    tar_mask = cv2.resize(tar_mask, size)
+
+    masked_ref_image = cv2.resize(masked_ref_image.astype(np.uint8), size).astype(np.uint8)
+    pipe_prior_output = redux(Image.fromarray(masked_ref_image))
+
+
+    tar_image = pad_to_square(tar_image, pad_value=255)
+
+    H2, W2 = tar_image.shape[0], tar_image.shape[1]
+
+    tar_image = cv2.resize(tar_image, size)
+    diptych_ref_tar = np.concatenate([masked_ref_image, tar_image], axis=1)
+
+
+    tar_mask = np.stack([tar_mask,tar_mask,tar_mask],-1)
+    mask_black = np.ones_like(tar_image) * 0
+    mask_diptych = np.concatenate([mask_black, tar_mask], axis=1)
+
+
+    diptych_ref_tar = Image.fromarray(diptych_ref_tar)
+    mask_diptych[mask_diptych == 1] = 255
+    mask_diptych = Image.fromarray(mask_diptych)
+
+
+
+    generator = torch.Generator("cuda").manual_seed(seed)
+    edited_image = pipe(
+        image=diptych_ref_tar,
+        mask_image=mask_diptych,
+        height=mask_diptych.size[1],
+        width=mask_diptych.size[0],
+        max_sequence_length=512,
+        generator=generator,
+        **pipe_prior_output, 
+    ).images[0]
+
+
+
+    width, height = edited_image.size
+    left = width // 2
+    right = width
+    top = 0
+    bottom = height
+    edited_image = edited_image.crop((left, top, right, bottom))
+
+
+    edited_image = np.array(edited_image)
+    edited_image = crop_back(edited_image, old_tar_image, np.array([H1, W1, H2, W2]), np.array(tar_box_yyxx_crop)) 
+    edited_image = Image.fromarray(edited_image)
+
+
+    return [edited_image]
+
+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("#  Play with InsertAnything to Insert your Target Objects! ")
+    gr.Markdown("# Upload / Draw Images for the Background (up) and Reference Object (down)")
+    gr.Markdown("### Draw mask on the background or just upload the mask.")
+    gr.Markdown("### Only select one of these two methods. Don't forget to click the corresponding button!!")
+
+    with gr.Row():
+        with gr.Column():
+            with gr.Row():
+                base_image = gr.Image(label="Background Image", source="upload", tool="sketch", type="pil",
+                                        brush_color='#FFFFFF', mask_opacity=0.5)
+
+                base_mask = gr.Image(label="Background Mask", source="upload", type="pil")
+
+            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.Image(label="Reference Image", source="upload", tool="sketch", type="pil", 
+                                    brush_color='#FFFFFF', mask_opacity=0.5)
+                                    
+                ref_mask = gr.Image(label="Reference Mask", source="upload", type="pil")
+
+            with gr.Row():
+                ref_mask_option = gr.Radio(["Draw Mask", "Upload with Mask"], label="Reference Mask Input Option", value="Upload with Mask")
+
+            baseline_gallery = gr.Gallery(label='Output', show_label=True, elem_id="gallery", height=512, columns=1)
+            with gr.Accordion("Advanced Option", open=True):
+                seed = gr.Slider(label="Seed", minimum=-1, maximum=999999999, 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.")
+
+    run_local_button = gr.Button(value="Run")
+
+
+    # #### example #####
+    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:
+            with gr.Row():
+                gr.HTML("<hr>")
+    # #### example #####
+        
+    run_local_button.click(fn=run_local, 
+                           inputs=[base_image, base_mask, ref_image, ref_mask, seed, base_mask_option, ref_mask_option], 
+                           outputs=[baseline_gallery]
+                        )
+demo.launch()

requirements.txt

@@ -0,0 +1,15 @@
+torch==2.5.1
+torchvision==0.20.1
+diffusers==0.32.2
+transformers==4.50.3
+peft==0.15.1
+opencv-python
+protobuf
+sentencepiece
+gradio==3.39.0
+bezier
+lightning==2.5.1
+datasets
+prodigyopt
+einops
+scipy

utils/utils.py

@@ -0,0 +1,140 @@
+import numpy as np
+import cv2
+import math
+
+def f(r, T=0.6, beta=0.1):
+    return np.where(r < T, beta + (1 - beta) / T * r, 1)
+
+# Get the bounding box of the mask
+def get_bbox_from_mask(mask):
+    h,w = mask.shape[0],mask.shape[1]
+
+    if mask.sum() < 10:
+        return 0,h,0,w
+    rows = np.any(mask,axis=1)
+    cols = np.any(mask,axis=0)
+    y1,y2 = np.where(rows)[0][[0,-1]]
+    x1,x2 = np.where(cols)[0][[0,-1]]
+    return (y1,y2,x1,x2)
+
+# Expand the bounding box
+def expand_bbox(mask, yyxx, ratio, min_crop=0):
+    y1,y2,x1,x2 = yyxx
+    H,W = mask.shape[0], mask.shape[1]
+
+    yyxx_area = (y2-y1+1) * (x2-x1+1)
+    r1 = yyxx_area / (H * W)
+    r2 = f(r1)
+    ratio = math.sqrt(r2 / r1)
+
+    xc, yc = 0.5 * (x1 + x2), 0.5 * (y1 + y2)
+    h = ratio * (y2-y1+1)
+    w = ratio * (x2-x1+1)
+    h = max(h,min_crop)
+    w = max(w,min_crop)
+
+    x1 = int(xc - w * 0.5)
+    x2 = int(xc + w * 0.5)
+    y1 = int(yc - h * 0.5)
+    y2 = int(yc + h * 0.5)
+
+    x1 = max(0,x1)
+    x2 = min(W,x2)
+    y1 = max(0,y1)
+    y2 = min(H,y2)
+    return (y1,y2,x1,x2)
+
+# Pad the image to a square shape
+def pad_to_square(image, pad_value = 255, random = False):
+    H,W = image.shape[0], image.shape[1]
+    if H == W:
+        return image
+
+    padd = abs(H - W)
+    if random:
+        padd_1 = int(np.random.randint(0,padd))
+    else:
+        padd_1 = int(padd / 2)
+    padd_2 = padd - padd_1
+
+    if len(image.shape) == 2: 
+        if H > W:
+            pad_param = ((0, 0), (padd_1, padd_2))
+        else:
+            pad_param = ((padd_1, padd_2), (0, 0))
+    elif len(image.shape) == 3: 
+        if H > W:
+            pad_param = ((0, 0), (padd_1, padd_2), (0, 0))
+        else:
+            pad_param = ((padd_1, padd_2), (0, 0), (0, 0))
+
+    image = np.pad(image, pad_param, 'constant', constant_values=pad_value)
+
+    return image
+
+# Expand the image and mask
+def expand_image_mask(image, mask, ratio=1.4):
+    h,w = image.shape[0], image.shape[1]
+    H,W = int(h * ratio), int(w * ratio) 
+    h1 = int((H - h) // 2)
+    h2 = H - h - h1
+    w1 = int((W -w) // 2)
+    w2 = W -w - w1
+
+    pad_param_image = ((h1,h2),(w1,w2),(0,0))
+    pad_param_mask = ((h1,h2),(w1,w2))
+    image = np.pad(image, pad_param_image, 'constant', constant_values=255)
+    mask = np.pad(mask, pad_param_mask, 'constant', constant_values=0)
+    return image, mask
+
+# Convert the bounding box to a square shape
+def box2squre(image, box):
+    H,W = image.shape[0], image.shape[1]
+    y1,y2,x1,x2 = box
+    cx = (x1 + x2) // 2
+    cy = (y1 + y2) // 2
+    h,w = y2-y1, x2-x1
+
+    if h >= w:
+        x1 = cx - h//2
+        x2 = cx + h//2
+    else:
+        y1 = cy - w//2
+        y2 = cy + w//2
+    x1 = max(0,x1)
+    x2 = min(W,x2)
+    y1 = max(0,y1)
+    y2 = min(H,y2)
+    return (y1,y2,x1,x2)
+
+# Crop the predicted image back to the original image
+def crop_back( pred, tar_image,  extra_sizes, tar_box_yyxx_crop):
+    H1, W1, H2, W2 = extra_sizes
+    y1,y2,x1,x2 = tar_box_yyxx_crop    
+    pred = cv2.resize(pred, (W2, H2))
+    m = 2 # maigin_pixel
+
+    if W1 == H1:
+        if m != 0:
+            tar_image[y1+m :y2-m, x1+m:x2-m, :] =  pred[m:-m, m:-m]
+        else:
+            tar_image[y1 :y2, x1:x2, :] =  pred[:, :]
+        return tar_image
+
+    if W1 < W2:
+        pad1 = int((W2 - W1) / 2)
+        pad2 = W2 - W1 - pad1
+        pred = pred[:,pad1: -pad2, :]
+    else:
+        pad1 = int((H2 - H1) / 2)
+        pad2 = H2 - H1 - pad1
+        pred = pred[pad1: -pad2, :, :]
+
+    gen_image = tar_image.copy()
+    if m != 0:
+        gen_image[y1+m :y2-m, x1+m:x2-m, :] =  pred[m:-m, m:-m]
+    else:
+        gen_image[y1 :y2, x1:x2, :] =  pred[:, :]
+    
+    return gen_image
+