app.py

import os

os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

import spaces  # noqa: E402  (must come before torch / CUDA-touching imports)
import math
import time
import random

import numpy as np
import torch
import gradio as gr
from PIL import Image

from direct import DirectPipeline

# ----------------------------------------------------------------------------
# Config
# ----------------------------------------------------------------------------
MODEL_INPUT_RESOLUTION = 1024
DIRECT_MODEL_PATH = "superGong/DIRECT"
FLUX_MODEL_PATH = "black-forest-labs/FLUX.1-Fill-dev"
SIGLIP_MODEL_PATH = "google/siglip2-so400m-patch14-384"

HF_TOKEN = os.environ.get("HF_TOKEN")

# ----------------------------------------------------------------------------
# Load models at module scope (ZeroGPU packs weights to disk after this)
# ----------------------------------------------------------------------------
print("Loading DIRECT pipeline (FLUX.1-Fill-dev + SigLIP2 + DIRECT adapters)...")
direct_pipeline = DirectPipeline.from_pretrained(
    direct_model_path=DIRECT_MODEL_PATH,
    flux_model_path=FLUX_MODEL_PATH,
    siglip_model_path=SIGLIP_MODEL_PATH,
    device=torch.device("cuda"),
    torch_dtype=torch.bfloat16,
    token=HF_TOKEN,
)
print("DIRECT pipeline loaded.")

# Background remover for the object image (ungated). Loaded lazily/cheaply.
_rembg_session = None


def _get_rembg_session():
    global _rembg_session
    if _rembg_session is None:
        from rembg import new_session

        _rembg_session = new_session("u2net")
    return _rembg_session


# ----------------------------------------------------------------------------
# Image-preparation helpers (2D proxy construction).
#
# The full DIRECT paper uses an interactive 3D viewer (TRELLIS + Viser) to let
# users pose a reconstructed 3D proxy of the object. That live 3D websocket
# viewer cannot run inside a single-port HF Space, so here we build the model's
# geometric-guidance inputs from a simple 2D placement (position + scale). The
# underlying DIRECT model (real weights) then performs the 3D-aware harmonized
# insertion. See the notes in the UI for this limitation.
# ----------------------------------------------------------------------------

def segment_object(object_rgb: Image.Image) -> Image.Image:
    """Return an RGBA image of the object with background removed."""
    from rembg import remove

    rgba = remove(object_rgb.convert("RGB"), session=_get_rembg_session())
    return rgba.convert("RGBA")


def _tight_crop_rgba(rgba: Image.Image) -> Image.Image:
    alpha = np.array(rgba.split()[-1])
    ys, xs = np.where(alpha > 10)
    if ys.size == 0:
        return rgba
    y1, y2, x1, x2 = ys.min(), ys.max() + 1, xs.min(), xs.max() + 1
    return rgba.crop((x1, y1, x2, y2))


def center_reference(rgba: Image.Image, out_size: int = MODEL_INPUT_RESOLUTION) -> Image.Image:
    """Object centered on black, square, with ~1.2 margin (model reference input)."""
    obj = _tight_crop_rgba(rgba)
    w, h = obj.size
    side = max(int(math.ceil(max(w, h) * 1.2)), 1)
    canvas = Image.new("RGB", (side, side), (0, 0, 0))
    canvas.paste(obj, ((side - w) // 2, (side - h) // 2), obj)
    return canvas.resize((out_size, out_size), Image.LANCZOS)


def place_object(bg: Image.Image, obj_rgba: Image.Image, cx: float, cy: float, scale: float):
    """Paste the (tight-cropped) object onto a copy of the background.

    cx, cy in [0, 1] (center), scale in [0, 1] (object longest side as a
    fraction of the background's longest side). Returns (placed_rgb, mask_L).
    """
    bg = bg.convert("RGB")
    W, H = bg.size
    obj = _tight_crop_rgba(obj_rgba)
    ow, oh = obj.size
    target_long = max(1, int(scale * max(W, H)))
    ratio = target_long / max(ow, oh)
    new_w = max(1, int(ow * ratio))
    new_h = max(1, int(oh * ratio))
    obj_r = obj.resize((new_w, new_h), Image.LANCZOS)

    center_x = int(cx * W)
    center_y = int(cy * H)
    x0 = center_x - new_w // 2
    y0 = center_y - new_h // 2

    placed_rgb = bg.copy()
    placed_rgb.paste(obj_r, (x0, y0), obj_r)

    mask = Image.new("L", (W, H), 0)
    obj_alpha = obj_r.split()[-1]
    mask.paste(obj_alpha, (x0, y0), obj_alpha)

    # Geometry proxy: the object RGB on a black canvas at its placed location.
    geometry_full = Image.new("RGB", (W, H), (0, 0, 0))
    geometry_full.paste(obj_r, (x0, y0), obj_r)

    return placed_rgb, mask, geometry_full


def get_mask_bbox(mask_pil, threshold=20):
    arr = np.array(mask_pil)
    ys, xs = np.where(arr > threshold)
    if ys.size == 0:
        return None
    return (xs.min(), ys.min(), xs.max() + 1, ys.max() + 1)


def get_smart_crop_bbox(mask_pil, min_ratio=0.02, max_ratio=0.3):
    bbox = get_mask_bbox(mask_pil)
    if bbox is None:
        s = MODEL_INPUT_RESOLUTION
        return (0, 0, s, s), s
    min_x, min_y, max_x, max_y = bbox
    mask_w, mask_h = max_x - min_x, max_y - min_y
    area = mask_w * mask_h
    side = int(math.sqrt(area / ((min_ratio + max_ratio) / 2.0)))
    side = max(side, max(mask_w, mask_h) + 40)
    cx = (min_x + max_x) // 2
    cy = (min_y + max_y) // 2
    half = side // 2
    return (cx - half, cy - half, cx - half + side, cy - half + side), side


def crop_and_pad(image, bbox, target_side):
    x1, y1, x2, y2 = bbox
    W, H = image.size
    valid = image.crop((max(0, x1), max(0, y1), min(W, x2), min(H, y2)))
    canvas = Image.new(image.mode, (target_side, target_side), 0)
    canvas.paste(valid, (max(0, -x1), max(0, -y1)))
    return canvas


def dilate_mask(mask_np, radius=10):
    import cv2

    m = (mask_np > 0).astype(np.uint8) * 255
    k = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (radius * 2 + 1, radius * 2 + 1))
    return (cv2.dilate(m, k, iterations=1) > 0).astype(np.uint8)


def refine_mask_holes(mask_bool, kernel_size=7):
    import cv2

    m = mask_bool.astype(np.uint8) * 255
    k = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (kernel_size, kernel_size))
    closed = cv2.morphologyEx(m, cv2.MORPH_CLOSE, k, iterations=2)
    contours, _ = cv2.findContours(closed, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    filled = np.zeros_like(closed)
    cv2.drawContours(filled, contours, -1, 255, thickness=cv2.FILLED)
    return filled > 127


def adain_color_fix(target_pil, source_pil, mask_pil):
    from torchvision.transforms import ToPILImage, ToTensor

    to_tensor = ToTensor()
    t = to_tensor(target_pil).unsqueeze(0)
    s = to_tensor(source_pil).unsqueeze(0)
    m = to_tensor(mask_pil).unsqueeze(0)
    eps = 1e-5
    res = t.clone()
    for ch in range(3):
        bg_idx = m[0, 0] < 0.1
        if bg_idx.sum() < 10:
            continue
        s_pix = s[0, ch][bg_idx]
        t_pix = t[0, ch][bg_idx]
        s_mean, s_std = s_pix.mean(), s_pix.std() + eps
        t_mean, t_std = t_pix.mean(), t_pix.std() + eps
        res[0, ch] = (t[0, ch] - t_mean) * (s_std / t_std) + s_mean
    return ToPILImage()(res.squeeze(0).clamp(0, 1))


def build_inputs(bg_pil, composite_full, mask_full, reference_ref, geometry_full):
    """Produce the model's 1024x1024 conditioning tensors from full-frame inputs."""
    target_res = MODEL_INPUT_RESOLUTION

    mask_np = np.array(mask_full)
    dilated01 = dilate_mask(mask_np, radius=10)
    dilated_pil = Image.fromarray(dilated01 * 255, mode="L")

    # Context image: full background with the (dilated) insertion region blacked.
    full_bg = np.array(bg_pil.convert("RGB"))
    context_image = Image.fromarray((full_bg * (1 - dilated01[:, :, None])).astype(np.uint8))

    ideal_bbox, target_side = get_smart_crop_bbox(dilated_pil)

    patch_composite = crop_and_pad(composite_full, ideal_bbox, target_side)
    patch_mask = crop_and_pad(dilated_pil, ideal_bbox, target_side)
    patch_geometry = crop_and_pad(geometry_full, ideal_bbox, target_side)
    patch_bg_ref = crop_and_pad(bg_pil.convert("RGB"), ideal_bbox, target_side)
    patch_mask_orig = crop_and_pad(Image.fromarray(mask_np), ideal_bbox, target_side)

    comp_arr = np.array(patch_composite)
    mask_dilated_arr = np.array(patch_mask) > 127
    mask_orig_arr = refine_mask_holes(np.array(patch_mask_orig) > 127, kernel_size=7)
    diff_region = mask_dilated_arr & (~mask_orig_arr)
    comp_arr[diff_region] = [0, 0, 0]
    patch_composite = Image.fromarray(comp_arr)

    composite_image = patch_composite.resize((target_res, target_res), Image.LANCZOS)
    model_input_mask = Image.fromarray(np.array(patch_mask).astype(np.uint8)).resize(
        (target_res, target_res), Image.NEAREST
    )
    geometry_image = patch_geometry.resize((target_res, target_res), Image.LANCZOS)
    background_reference_image = patch_bg_ref.resize((target_res, target_res), Image.LANCZOS)

    inpaint_mask = Image.fromarray(((np.array(model_input_mask) > 0) * 255).astype(np.uint8))

    return {
        "composite_image": composite_image,
        "inpaint_mask": inpaint_mask,
        "reference_image": reference_ref,
        "geometry_image": geometry_image,
        "context_image": context_image,
        "model_input_mask": model_input_mask,
        "background_reference_image": background_reference_image,
        "ideal_bbox": ideal_bbox,
        "target_side": target_side,
    }


def paste_back(bg_pil, generated_patch, inp):
    fixed = adain_color_fix(
        generated_patch, inp["background_reference_image"], inp["model_input_mask"]
    )
    fixed = fixed.resize((inp["target_side"], inp["target_side"]), Image.LANCZOS)
    x1, y1, x2, y2 = inp["ideal_bbox"]
    W, H = bg_pil.size
    pad_left = max(0, -x1)
    pad_top = max(0, -y1)
    valid_w = min(W, x2) - max(0, x1)
    valid_h = min(H, y2) - max(0, y1)
    patch_valid = fixed.crop((pad_left, pad_top, pad_left + valid_w, pad_top + valid_h))
    out = bg_pil.convert("RGB").copy()
    out.paste(patch_valid, (max(0, x1), max(0, y1)))
    return out


# ----------------------------------------------------------------------------
# Inference
# ----------------------------------------------------------------------------

def _estimate_duration(bg, obj, cx, cy, scale, seed, ref_scale, steps, *a, **k):
    # Measured ~12 s/step at 1024 when reference guidance is on (CFG doubles the
    # forward pass); ~half that when it is off. Plus fixed overhead for VAE /
    # rembg / cold worker init.
    try:
        steps = int(steps)
    except Exception:
        steps = 16
    try:
        ref_on = float(ref_scale) > 1.0
    except Exception:
        ref_on = True
    per_step = 12.5 if ref_on else 6.5
    return int(min(600, 45 + steps * per_step))


@spaces.GPU(duration=_estimate_duration)
def insert_object(
    bg: Image.Image,
    obj: Image.Image,
    cx: float,
    cy: float,
    scale: float,
    seed: int,
    ref_scale: float,
    steps: int,
    progress=gr.Progress(track_tqdm=True),
):
    """Insert a reference object into a background image with 3D-aware harmonization.

    Args:
        bg: Background scene image.
        obj: Reference object image (background is removed automatically).
        cx: Horizontal placement of the object center (0=left, 1=right).
        cy: Vertical placement of the object center (0=top, 1=bottom).
        scale: Object size as a fraction of the background's longest side.
        seed: Random seed for reproducibility.
        ref_scale: Reference guidance scale (identity preservation strength).
        steps: Number of inference steps.

    Returns:
        The composited image with the object inserted, and a preview of the raw
        2D placement used as geometric guidance.
    """
    if bg is None:
        raise gr.Error("Please provide a background image.")
    if obj is None:
        raise gr.Error("Please provide an object image.")

    t0 = time.perf_counter()
    bg = bg.convert("RGB")
    obj_rgba = segment_object(obj)

    reference_ref = center_reference(obj_rgba, out_size=MODEL_INPUT_RESOLUTION)
    placed_rgb, mask_full, geometry_full = place_object(bg, obj_rgba, cx, cy, scale)

    inp = build_inputs(bg, placed_rgb, mask_full, reference_ref, geometry_full)

    seed = int(seed)
    final_images = direct_pipeline(
        composite_image=inp["composite_image"],
        inpaint_mask=inp["inpaint_mask"],
        reference_image=inp["reference_image"],
        geometry_image=inp["geometry_image"],
        context_image=inp["context_image"],
        seed=seed,
        guidance_scale=30,
        num_inference_steps=int(steps),
        height=MODEL_INPUT_RESOLUTION,
        width=MODEL_INPUT_RESOLUTION,
        use_autocast=True,
        reference_guidance_scale=float(ref_scale),
    )
    generated_patch = final_images[0]
    result = paste_back(bg, generated_patch, inp)
    print(f"[insert_object] done in {time.perf_counter() - t0:.1f}s (steps={steps})")
    return result, placed_rgb


def randomize_seed():
    return random.randint(0, 2**31 - 1)


# ----------------------------------------------------------------------------
# UI
# ----------------------------------------------------------------------------
CSS = """
#col-container { max-width: 1200px; margin: 0 auto; }
.dark .gradio-container { color: var(--body-text-color); }
"""

INTRO = """
# DIRECT: 3D-Aware Object Insertion

Insert a reference **object** into a **background** scene with realistic,
harmonized results, powered by the [DIRECT](https://huggingface.co/superGong/DIRECT)
model (ICML 2026) — a FLUX.1-Fill-dev network guided by a decomposed visual proxy.

**How to use:** upload a background and an object image (its background is
removed automatically), choose *where* and *how big* to place it, then click **Insert**.

> **Note.** The full paper uses an interactive 3D viewer (TRELLIS + Viser) to pose a
> reconstructed 3D proxy of the object. That live 3D viewer cannot run inside a
> single-port Space, so this demo drives the same DIRECT model with a simpler
> **2D placement** (position + scale) as its geometric guidance.

[Paper](https://arxiv.org/abs/2606.06601) · [Project page](https://gong1130.github.io/DIRECT/) · [Code](https://github.com/Gong1130/DIRECT)
"""

with gr.Blocks(theme=gr.themes.Citrus(), css=CSS) as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown(INTRO)
        with gr.Row():
            with gr.Column(scale=1):
                bg_input = gr.Image(label="Background image", type="pil", height=300)
                obj_input = gr.Image(label="Object image", type="pil", height=300)
                run_btn = gr.Button("Insert", variant="primary")
            with gr.Column(scale=1):
                out_result = gr.Image(label="Inserted result", type="pil", height=360)
                out_preview = gr.Image(label="2D placement (geometric guidance)", type="pil", height=240)

        with gr.Accordion("Placement & advanced settings", open=True):
            with gr.Row():
                cx = gr.Slider(0.0, 1.0, value=0.5, step=0.01, label="Horizontal position")
                cy = gr.Slider(0.0, 1.0, value=0.6, step=0.01, label="Vertical position")
                scale = gr.Slider(0.05, 0.9, value=0.35, step=0.01, label="Object size")
            with gr.Row():
                ref_scale = gr.Slider(1.0, 5.0, value=2.0, step=0.1, label="Reference guidance scale")
                steps = gr.Slider(12, 28, value=16, step=1, label="Inference steps")
                seed = gr.Number(label="Seed", value=42, precision=0)
            rand_btn = gr.Button("🎲 Randomize seed")

        gr.Examples(
            examples=[
                ["examples/bg_landscape.jpg", "examples/obj_ducks.jpg", 0.55, 0.70, 0.28, 42, 2.0, 16],
                ["examples/bg_tent.jpg", "examples/obj_dog.jpg", 0.45, 0.68, 0.30, 7, 2.0, 16],
                ["examples/bg_beach.jpg", "examples/obj_cake.jpg", 0.50, 0.72, 0.22, 123, 2.5, 16],
            ],
            inputs=[bg_input, obj_input, cx, cy, scale, seed, ref_scale, steps],
            outputs=[out_result, out_preview],
            fn=insert_object,
            cache_examples=True,
            cache_mode="lazy",
        )

    rand_btn.click(fn=randomize_seed, outputs=seed)
    run_btn.click(
        fn=insert_object,
        inputs=[bg_input, obj_input, cx, cy, scale, seed, ref_scale, steps],
        outputs=[out_result, out_preview],
        api_name="insert",
    )

if __name__ == "__main__":
    demo.launch(mcp_server=True)