Upload 18 files

.gitattributes

@@ -33,3 +33,11 @@ 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/background_image/1.png filter=lfs diff=lfs merge=lfs -text
+examples/background_image/2.png filter=lfs diff=lfs merge=lfs -text
+examples/background_image/3.png filter=lfs diff=lfs merge=lfs -text
+examples/background_image/4.png 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

app.py

@@ -0,0 +1,561 @@
+import base64
+from io import BytesIO
+
+import gradio as gr
+import torch
+from PIL import Image, ImageChops
+import spaces
+from diffsynth.pipelines.flux2_image import Flux2ImagePipeline, ModelConfig
+
+from utils import process_source, process_reference, paste_back, binarize_mask
+
+pipe = Flux2ImagePipeline.from_pretrained(
+    torch_dtype=torch.bfloat16,
+    device="cuda",
+    model_configs=[
+        ModelConfig(model_id="black-forest-labs/FLUX.2-klein-base-4B", origin_file_pattern="text_encoder/*.safetensors"),
+        ModelConfig(model_id="black-forest-labs/FLUX.2-klein-base-4B", origin_file_pattern="transformer/*.safetensors"),
+        ModelConfig(model_id="black-forest-labs/FLUX.2-klein-base-4B", origin_file_pattern="vae/diffusion_pytorch_model.safetensors"),
+    ],
+    tokenizer_config=ModelConfig(model_id="black-forest-labs/FLUX.2-klein-base-4B", origin_file_pattern="tokenizer/"),
+)
+
+pipe.load_lora(pipe.dit, ModelConfig(model_id="HuanJue/Easy-Insert", origin_file_pattern="easy-insert.safetensors"))
+
+PROMPT = "Replace the white mask of image1 with the content in image2. Preserving the background, lighting, and surrounding elements, maintain a seamless and natural result."
+
+
+# ===================== Example Data =====================
+
+CANVAS_W, CANVAS_H = 1024, 1024
+
+EXAMPLES = [
+    ("examples/background_image/1.png", "examples/insert_mask/1.png", "examples/ref_image/1.png", "examples/ref_mask/1.png"),
+    ("examples/background_image/2.png", "examples/insert_mask/2.png", "examples/ref_image/2.png", "examples/ref_mask/2.png"),
+    ("examples/background_image/3.png", "examples/insert_mask/3.png", "examples/ref_image/3.png", "examples/ref_mask/3.png"),
+    ("examples/background_image/4.png", "examples/insert_mask/4.png", "examples/ref_image/4.png", "examples/ref_mask/4.png"),
+]
+
+
+# ===================== Thumbnail HTML Generation =====================
+
+def img_to_b64(path, size=(240, 240)):
+    img = Image.open(path).convert("RGB")
+    # Preserve the original aspect ratio; fit inside `size`. 2x the display
+    # width (110 CSS px) so thumbnails stay sharp on HiDPI screens.
+    img.thumbnail(size, Image.LANCZOS)
+    buf = BytesIO()
+    img.save(buf, format="PNG")
+    return base64.b64encode(buf.getvalue()).decode()
+
+
+def build_examples_html():
+    row_pairs = [EXAMPLES[0:2], EXAMPLES[2:4]]
+
+    html = '<div class="ex-grid">'
+    for row_idx, row_examples in enumerate(row_pairs):
+        html += '<div class="ex-grid-row">'
+        for col_idx, (bg, mask, ref, ref_mask) in enumerate(row_examples):
+            i = row_idx * 2 + col_idx
+            bg_b64 = img_to_b64(bg)
+            mask_b64 = img_to_b64(mask)
+            ref_b64 = img_to_b64(ref)
+            ref_mask_b64 = img_to_b64(ref_mask)
+            html += f'''
+            <div class="ex-row" onclick="(function(){{var el=document.getElementById('ex_btn_{i}');if(!el)return;var btn=el.querySelector('button')||el;btn.dispatchEvent(new MouseEvent('click',{{bubbles:true,cancelable:true}}));}})()">
+                <div class="ex-label">Example {i + 1}</div>
+                <div class="ex-thumbs">
+                    <div class="ex-thumb-wrap">
+                        <img src="data:image/png;base64,{bg_b64}" class="ex-thumb" draggable="false"/>
+                        <span class="ex-thumb-sublabel">Background</span>
+                    </div>
+                    <div class="ex-thumb-wrap">
+                        <img src="data:image/png;base64,{mask_b64}" class="ex-thumb" draggable="false"/>
+                        <span class="ex-thumb-sublabel">Mask</span>
+                    </div>
+                    <div class="ex-thumb-wrap">
+                        <img src="data:image/png;base64,{ref_b64}" class="ex-thumb" draggable="false"/>
+                        <span class="ex-thumb-sublabel">Reference</span>
+                    </div>
+                    <div class="ex-thumb-wrap">
+                        <img src="data:image/png;base64,{ref_mask_b64}" class="ex-thumb" draggable="false"/>
+                        <span class="ex-thumb-sublabel">Ref Mask</span>
+                    </div>
+                </div>
+            </div>
+            '''
+        html += '</div>'
+    html += '</div>'
+    return html
+
+
+# ===================== Utility Functions =====================
+
+def extract_mask_from_layers(layers, target_size):
+    mask = Image.new("L", target_size, 0)
+    for layer in layers:
+        if layer is not None:
+            layer_rgba = layer.convert("RGBA").resize(target_size)
+            alpha = layer_rgba.split()[3]
+            alpha_binary = alpha.point(lambda x: 255 if x > 0 else 0)
+            mask = Image.composite(Image.new("L", target_size, 255), mask, alpha_binary)
+    return mask
+
+
+EDITOR_SIZE = 800
+# Matches Gradio light-mode image editor background so padded examples blend in.
+EDITOR_BG_COLOR = (243, 244, 246)
+
+
+def fit_transform(img, editor_size=EDITOR_SIZE):
+    """Display transform the fixed-canvas editor applies to `img`: scale to fit
+    inside the square canvas, then center-pad. Returns (disp_size, pad)."""
+    scale = min(editor_size / img.width, editor_size / img.height, 1.0)
+    disp_size = (
+        max(1, round(img.width * scale)),
+        max(1, round(img.height * scale)),
+    )
+    pad = (
+        (editor_size - disp_size[0]) // 2,
+        (editor_size - disp_size[1]) // 2,
+    )
+    return disp_size, pad
+
+
+def build_display_bg(bg, editor_size=EDITOR_SIZE):
+    """Rebuild the padded square background the editor shows for `bg`."""
+    disp_size, pad = fit_transform(bg, editor_size)
+    canvas = Image.new("RGB", (editor_size, editor_size), EDITOR_BG_COLOR)
+    canvas.paste(bg.convert("RGB").resize(disp_size, Image.LANCZOS), pad)
+    return canvas
+
+
+def images_close(a, b, tol_frac=0.001):
+    """True if two RGB images are (nearly) pixel-identical."""
+    if a.size != b.size:
+        return False
+    diff = ImageChops.difference(a.convert("RGB"), b.convert("RGB"))
+    changed = sum(diff.histogram()[1:])
+    return changed <= tol_frac * a.size[0] * a.size[1] * 3
+
+
+def make_editor_value(bg_pil, mask_pil=None, editor_size=EDITOR_SIZE):
+    """Build an ImageEditor value on a fixed square canvas.
+
+    The background is downscaled to fit inside `editor_size` and padded to a
+    square; the mask is padded the same way and sent as a real layer. This
+    keeps the mask and background perfectly aligned in the frontend (both are
+    800x800 and the padding is identical), while the full-resolution image and
+    exact mask are kept in gr.State for inference.
+
+    Returns (editor_value, sent_mask, transform).
+    """
+    bg = bg_pil.convert("RGB")
+    disp_size, pad = fit_transform(bg, editor_size)
+    bg_sq = build_display_bg(bg, editor_size)
+
+    transform = {
+        "editor_size": editor_size,
+        "disp_size": disp_size,
+        "pad": pad,
+    }
+
+    if mask_pil is None:
+        # composite=None avoids saving a second, redundant image file. The editor
+        # can render the background directly; if it needs a composite it will build
+        # it from background + (empty) layers.
+        return {"background": bg_sq, "layers": [], "composite": None}, None, transform
+
+    mask = mask_pil.convert("L").resize(disp_size, Image.NEAREST)
+    mask_sq = Image.new("L", (editor_size, editor_size), 0)
+    mask_sq.paste(mask, pad)
+
+    transparent = Image.new("RGBA", (editor_size, editor_size), (0, 0, 0, 0))
+    # Semi-transparent white (alpha 153 ≈ 60%) so the background shows through
+    # the loaded mask. Mask extraction binarizes alpha > 0, so this still
+    # resolves to a solid mask at inference time.
+    white_solid = Image.new("RGBA", (editor_size, editor_size), (255, 255, 255, 153))
+    mask_layer = Image.composite(white_solid, transparent, mask_sq)
+
+    return {"background": bg_sq, "layers": [mask_layer], "composite": None}, mask_sq, transform
+
+
+def masks_close(a, b, tol_frac=0.0):
+    """True if two binarized masks are pixel-identical (within tol_frac).
+    The fixed-canvas PNG round-trip is exact, so we treat any non-zero
+    difference as a real user edit."""
+    if a.size != b.size:
+        b = b.resize(a.size, Image.NEAREST)
+    diff = ImageChops.difference(binarize_mask(a), binarize_mask(b))
+    changed = sum(diff.histogram()[1:])
+    return changed <= tol_frac * a.size[0] * a.size[1]
+
+
+def load_example(idx):
+    bg_path, mask_path, ref_path, ref_mask_path = EXAMPLES[idx]
+    bg_img = Image.open(bg_path).convert("RGB")
+    ref_img = Image.open(ref_path).convert("RGB")
+    base_val, base_sent, base_transform = make_editor_value(
+        bg_img, Image.open(mask_path)
+    )
+    ref_val, ref_sent, ref_transform = make_editor_value(
+        ref_img, Image.open(ref_mask_path)
+    )
+    # Keep the pristine full-res image + exact mask in State. The editor value
+    # is only for display; its small size avoids slow frontend re-uploads.
+    base_state = (bg_img, Image.open(mask_path).convert("L"), base_sent, base_transform)
+    ref_state = (ref_img, Image.open(ref_mask_path).convert("L"), ref_sent, ref_transform)
+    return base_val, ref_val, base_state, ref_state
+
+
+def load_ex1():
+    return load_example(0)
+
+def load_ex2():
+    return load_example(1)
+
+def load_ex3():
+    return load_example(2)
+
+def load_ex4():
+    return load_example(3)
+
+
+# ===================== Generation Function =====================
+
+def resolve_source(editor_value, state):
+    """Return (full_res_image, full_res_mask) for one editor.
+
+    `state` is (image, mask, sent_mask, transform) captured when an example was
+    loaded. The editor value's background is display-only, so we don't trust it
+    blindly: rebuild the padded display image from the state's pristine image
+    and compare it with what the editor currently shows. If they match, the
+    example is still loaded and the state's full-res image (+ exact mask) is
+    used. If they differ, the user uploaded a different image - the editor
+    background is then the full-res original, and the brush strokes are mapped
+    back through the fixed-canvas transform.
+    """
+    bg_ed = editor_value.get("background")
+    if bg_ed is None:
+        return None, None
+    img_ed = bg_ed.convert("RGB")
+
+    ed_mask = extract_mask_from_layers(
+        editor_value.get("layers", []), (EDITOR_SIZE, EDITOR_SIZE)
+    )
+
+    if state is not None:
+        img, full_mask, sent_mask, transform = state
+        img = img.convert("RGB")
+        if images_close(img_ed, build_display_bg(img)):
+            if full_mask is not None:
+                # Example flow: exact full-res mask is in state. If the editor's
+                # layer mask still matches what we sent (user hasn't brushed or
+                # erased), use the state mask directly - bit-exact with inference.py.
+                if sent_mask is not None and masks_close(ed_mask, sent_mask):
+                    return img, binarize_mask(full_mask)
+                # User edited the example's mask (erased and/or brushed): crop the
+                # padding and map the edited mask back to the original resolution.
+                if transform:
+                    ed_mask_bin = binarize_mask(ed_mask)
+                    px, py = transform["pad"]
+                    dw, dh = transform["disp_size"]
+                    cropped = ed_mask_bin.crop((px, py, px + dw, py + dh))
+                    return img, binarize_mask(cropped.resize(img.size, Image.NEAREST))
+
+    # Upload flow: the editor background is the user's full-res original.
+    disp_size, pad = fit_transform(img_ed)
+    ed_mask_bin = binarize_mask(ed_mask)
+    cropped = ed_mask_bin.crop((pad[0], pad[1], pad[0] + disp_size[0], pad[1] + disp_size[1]))
+    return img_ed, binarize_mask(cropped.resize(img_ed.size, Image.NEAREST))
+
+@spaces.GPU
+def run_local(base, ref, base_state, ref_state, seed, num_inference_steps, cfg_scale):
+    if base is None or ref is None or not isinstance(base, dict) or not isinstance(ref, dict):
+        return None, gr.update(visible=False)
+
+    pil_bg, pil_mask = resolve_source(base, base_state)
+    pil_ref, pil_ref_mask = resolve_source(ref, ref_state)
+    if pil_bg is None or pil_ref is None:
+        return None, gr.update(visible=False)
+
+    if pil_mask.getextrema() == (0, 0) or pil_ref_mask.getextrema() == (0, 0):
+        error_html = """
+            <div class="error-overlay" style="
+                position: fixed; top: 0; left: 0; width: 100%; height: 100%;
+                background: rgba(0,0,0,0.5); display: flex; justify-content: center;
+                align-items: center; z-index: 9999;
+            ">
+                <div style="
+                    background: white; padding: 30px; border-radius: 10px;
+                    text-align: center; font-size: 18px; box-shadow: 0 0 15px rgba(0,0,0,0.3);
+                ">
+                    <p style="color: red; margin-bottom: 20px;">
+                        ⚠️ Please paint a mask on both the background image and the reference image first, or click an example!
+                    </p>
+                    <button onclick="this.closest('.error-overlay').remove()"
+                            style="padding: 8px 20px; cursor: pointer; border: none;
+                                background: #eee; border-radius: 5px;">
+                        OK
+                    </button>
+                </div>
+            </div>
+            """
+        return None, gr.update(value=error_html, visible=True)
+
+    background_image, _, crop_box, source_mask_cropped = process_source(pil_bg, pil_mask, CANVAS_W)
+    ref_image = process_reference(pil_ref, pil_ref_mask, CANVAS_W)
+
+    generated_image = pipe(
+        PROMPT,
+        edit_image=[background_image, ref_image],
+        seed=int(seed),
+        rand_device="cuda",
+        num_inference_steps=int(num_inference_steps),
+        cfg_scale=float(cfg_scale),
+        height=CANVAS_W,
+        width=CANVAS_H,
+    )
+    result_img = paste_back(generated_image, pil_bg, crop_box, source_mask_cropped, feather=0)
+
+    return result_img, gr.update(visible=False)
+
+
+# ===================== Gradio UI =====================
+
+with gr.Blocks(css="""
+    .input-row {
+        overflow: visible !important;
+    }
+
+    .input-row .gr-image-editor {
+        overflow: hidden !important;
+    }
+    .input-row .gr-image-editor .image-container,
+    .input-row .gr-image-editor .canvas-container,
+    .input-row .gr-image-editor canvas {
+        max-width: 100% !important;
+        max-height: 100% !important;
+        object-fit: contain !important;
+    }
+
+    .ex-section-header {
+        display: flex;
+        align-items: center;
+        gap: 10px;
+        margin: 28px 0 14px 0;
+        justify-content: center;
+    }
+    .ex-section-header::before {
+        content: '';
+        flex: 1;
+        height: 1px;
+        max-width: 180px;
+        background: #e5e7eb;
+    }
+    .ex-section-header::after {
+        content: '';
+        flex: 1;
+        height: 1px;
+        max-width: 180px;
+        background: #e5e7eb;
+    }
+
+    .ex-container {
+        display: flex;
+        flex-direction: column;
+        align-items: center;
+        padding-bottom: 20px;
+    }
+
+    .ex-grid {
+        display: flex;
+        flex-direction: column;
+        align-items: center;
+        gap: 12px;
+        padding-bottom: 20px;
+    }
+    .ex-grid-row {
+        display: flex;
+        gap: 20px;
+        justify-content: center;
+        flex-wrap: wrap;
+    }
+
+    .ex-row {
+        display: flex;
+        align-items: center;
+        gap: 20px;
+        padding: 14px 28px;
+        border: 2px solid #e5e7eb;
+        border-radius: 12px;
+        cursor: pointer;
+        transition: all 0.25s ease;
+        background: #ffffff;
+        user-select: none;
+        width: fit-content;
+    }
+    .ex-row:hover {
+        border-color: #3b82f6;
+        background: #f0f7ff;
+        box-shadow: 0 4px 18px rgba(59, 130, 246, 0.15);
+        transform: translateY(-2px);
+    }
+    .ex-row:active {
+        transform: translateY(0);
+        box-shadow: 0 2px 8px rgba(59, 130, 246, 0.2);
+    }
+
+    .ex-label {
+        font-weight: 700;
+        font-size: 15px;
+        min-width: 62px;
+        color: #1e40af;
+        letter-spacing: 0.02em;
+    }
+
+    .ex-thumbs {
+        display: flex;
+        gap: 14px;
+    }
+
+    .ex-thumb-wrap {
+        display: flex;
+        flex-direction: column;
+        align-items: center;
+        gap: 6px;
+    }
+
+    .ex-thumb {
+        width: 110px;
+        height: auto;
+        object-fit: contain;
+        border-radius: 8px;
+        border: 2px solid #e5e7eb;
+        transition: all 0.25s ease;
+        pointer-events: none;
+    }
+    .ex-row:hover .ex-thumb {
+        border-color: #93c5fd;
+    }
+
+    .ex-thumb-sublabel {
+        font-size: 12px;
+        color: #6b7280;
+        font-weight: 500;
+    }
+
+    /* Gradio >=5 does not mount components with visible=False into the DOM,
+       so the example-thumbnail JS cannot find them. Hide via CSS instead:
+       the buttons stay mounted (click handlers still fire) but invisible. */
+    .ex-hidden-btn {
+        display: none !important;
+    }
+""") as demo:
+
+    gr.Markdown(
+        "<h1 style='text-align: center;'>Reference-Based Object Insertion & Clothing Replacement</h1>"
+        "<h3 style='text-align: center;'>Insert an object from the reference image into the masked region of the background image, or replace the clothing in the masked region with the reference</h3>"
+        "<p style='text-align: center;'>"
+        "<a href='https://github.com/huan-yin/Easy-Insert' target='_blank' "
+        "style='color: #3b82f6; text-decoration: none; font-weight: 500;'>"
+        "GitHub Repo: huan-yin/Easy-Insert</a></p>"
+    )
+    gr.Markdown(
+        """
+        **Instructions:**
+        1. Upload a background image and use the brush/eraser tools below it to mark or modify the insertion region (mask); upload a reference image and paint over the object you want to insert (reference mask).
+        2. Or click any thumbnail row in the "Examples" section below to automatically load the background image + mask + reference image + reference mask.
+        3. Click the "Generate" button; the result will be shown below.
+        """
+    )
+
+    with gr.Row(elem_classes="input-row"):
+        base = gr.ImageEditor(
+            label="Background Image (paint the insertion region)",
+            type="pil",
+            format="png",
+            width=420,
+            height=450,
+            sources=["upload"],
+            canvas_size=(EDITOR_SIZE, EDITOR_SIZE),
+            fixed_canvas=True,
+            brush=gr.Brush(
+                default_size=30,
+                default_color="rgba(255, 255, 255, 0.6)",
+                color_mode="fixed",
+                colors=["rgba(255, 255, 255, 0.6)"],
+            ),
+            eraser=gr.Eraser(default_size=30),
+        )
+        ref = gr.ImageEditor(
+            label="Reference Image (paint the object to insert)",
+            type="pil",
+            format="png",
+            width=420,
+            height=450,
+            sources=["upload"],
+            canvas_size=(EDITOR_SIZE, EDITOR_SIZE),
+            fixed_canvas=True,
+            brush=gr.Brush(
+                default_size=30,
+                default_color="rgba(255, 255, 255, 0.6)",
+                color_mode="fixed",
+                colors=["rgba(255, 255, 255, 0.6)"],
+            ),
+            eraser=gr.Eraser(default_size=30),
+        )
+
+    with gr.Row():
+        seed = gr.Number(label="Random Seed", value=1, precision=0)
+        num_inference_steps = gr.Slider(label="Inference Steps", minimum=1, maximum=50, step=1, value=15)
+        cfg_scale = gr.Slider(label="CFG Scale", minimum=1, maximum=10, step=0.1, value=4)
+
+    with gr.Row():
+        gen_btn = gr.Button("Generate", variant="primary")
+
+    # ==================== Generation Result ====================
+    with gr.Row():
+        output_image = gr.Image(
+            label="Generated Result",
+            interactive=False,
+            width=420,
+            height=450,
+        )
+
+    with gr.Row():
+        error_dialog = gr.HTML(visible=False)
+
+    gr.HTML('<div class="ex-section-header"><span style="font-weight:700;font-size:16px;color:#374151;">Examples (click to load background image + mask + reference image + reference mask)</span></div>')
+
+    base_state = gr.State(None)
+    ref_state = gr.State(None)
+
+    ex_btn0 = gr.Button("Example 1", elem_id="ex_btn_0", elem_classes=["ex-hidden-btn"])
+    ex_btn1 = gr.Button("Example 2", elem_id="ex_btn_1", elem_classes=["ex-hidden-btn"])
+    ex_btn2 = gr.Button("Example 3", elem_id="ex_btn_2", elem_classes=["ex-hidden-btn"])
+    ex_btn3 = gr.Button("Example 4", elem_id="ex_btn_3", elem_classes=["ex-hidden-btn"])
+
+    gr.HTML('<div class="ex-container">' + build_examples_html() + '</div>')
+
+    # ==================== Event Bindings ====================
+
+    # Single-step load: a clear-then-load chain sends the frontend two rapid
+    # updates, and the upload round-trip in between references temp PNGs that
+    # get overwritten mid-read (truncated/broken PNG errors on preprocess).
+    ex_btn0.click(fn=load_ex1, outputs=[base, ref, base_state, ref_state])
+    ex_btn1.click(fn=load_ex2, outputs=[base, ref, base_state, ref_state])
+    ex_btn2.click(fn=load_ex3, outputs=[base, ref, base_state, ref_state])
+    ex_btn3.click(fn=load_ex4, outputs=[base, ref, base_state, ref_state])
+
+    gen_btn.click(
+        fn=run_local,
+        inputs=[base, ref, base_state, ref_state, seed, num_inference_steps, cfg_scale],
+        outputs=[output_image, error_dialog],
+    )
+
+demo.launch(server_name="0.0.0.0", server_port=7860)
+
+
+
+

requirements.txt

@@ -0,0 +1,10 @@
+diffsynth==2.1.2
+transformers
+accelerate
+peft
+safetensors
+sentencepiece
+torchvision
+Pillow
+numpy
+gradio==6.24.0