hugging-apps/direct-object-insertion
0
1import os2 3os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")4 5import spaces # noqa: E402 (must come before torch / CUDA-touching imports)6import math7import time8import random9 10import numpy as np11import torch12import gradio as gr13from PIL import Image14 15from direct import DirectPipeline16 17# ----------------------------------------------------------------------------18# Config19# ----------------------------------------------------------------------------20MODEL_INPUT_RESOLUTION = 102421DIRECT_MODEL_PATH = "superGong/DIRECT"22FLUX_MODEL_PATH = "black-forest-labs/FLUX.1-Fill-dev"23SIGLIP_MODEL_PATH = "google/siglip2-so400m-patch14-384"24 25HF_TOKEN = os.environ.get("HF_TOKEN")26 27# ----------------------------------------------------------------------------28# Load models at module scope (ZeroGPU packs weights to disk after this)29# ----------------------------------------------------------------------------30print("Loading DIRECT pipeline (FLUX.1-Fill-dev + SigLIP2 + DIRECT adapters)...")31direct_pipeline = DirectPipeline.from_pretrained(32 direct_model_path=DIRECT_MODEL_PATH,33 flux_model_path=FLUX_MODEL_PATH,34 siglip_model_path=SIGLIP_MODEL_PATH,35 device=torch.device("cuda"),36 torch_dtype=torch.bfloat16,37 token=HF_TOKEN,38)39print("DIRECT pipeline loaded.")40 41# Background remover for the object image (ungated). Loaded lazily/cheaply.42_rembg_session = None43 44 45def _get_rembg_session():46 global _rembg_session47 if _rembg_session is None:48 from rembg import new_session49 50 _rembg_session = new_session("u2net")51 return _rembg_session52 53 54# ----------------------------------------------------------------------------55# Image-preparation helpers (2D proxy construction).56#57# The full DIRECT paper uses an interactive 3D viewer (TRELLIS + Viser) to let58# users pose a reconstructed 3D proxy of the object. That live 3D websocket59# viewer cannot run inside a single-port HF Space, so here we build the model's60# geometric-guidance inputs from a simple 2D placement (position + scale). The61# underlying DIRECT model (real weights) then performs the 3D-aware harmonized62# insertion. See the notes in the UI for this limitation.63# ----------------------------------------------------------------------------64 65def segment_object(object_rgb: Image.Image) -> Image.Image:66 """Return an RGBA image of the object with background removed."""67 from rembg import remove68 69 rgba = remove(object_rgb.convert("RGB"), session=_get_rembg_session())70 return rgba.convert("RGBA")71 72 73def _tight_crop_rgba(rgba: Image.Image) -> Image.Image:74 alpha = np.array(rgba.split()[-1])75 ys, xs = np.where(alpha > 10)76 if ys.size == 0:77 return rgba78 y1, y2, x1, x2 = ys.min(), ys.max() + 1, xs.min(), xs.max() + 179 return rgba.crop((x1, y1, x2, y2))80 81 82def center_reference(rgba: Image.Image, out_size: int = MODEL_INPUT_RESOLUTION) -> Image.Image:83 """Object centered on black, square, with ~1.2 margin (model reference input)."""84 obj = _tight_crop_rgba(rgba)85 w, h = obj.size86 side = max(int(math.ceil(max(w, h) * 1.2)), 1)87 canvas = Image.new("RGB", (side, side), (0, 0, 0))88 canvas.paste(obj, ((side - w) // 2, (side - h) // 2), obj)89 return canvas.resize((out_size, out_size), Image.LANCZOS)90 91 92def place_object(bg: Image.Image, obj_rgba: Image.Image, cx: float, cy: float, scale: float):93 """Paste the (tight-cropped) object onto a copy of the background.94 95 cx, cy in [0, 1] (center), scale in [0, 1] (object longest side as a96 fraction of the background's longest side). Returns (placed_rgb, mask_L).97 """98 bg = bg.convert("RGB")99 W, H = bg.size100 obj = _tight_crop_rgba(obj_rgba)101 ow, oh = obj.size102 target_long = max(1, int(scale * max(W, H)))103 ratio = target_long / max(ow, oh)104 new_w = max(1, int(ow * ratio))105 new_h = max(1, int(oh * ratio))106 obj_r = obj.resize((new_w, new_h), Image.LANCZOS)107 108 center_x = int(cx * W)109 center_y = int(cy * H)110 x0 = center_x - new_w // 2111 y0 = center_y - new_h // 2112 113 placed_rgb = bg.copy()114 placed_rgb.paste(obj_r, (x0, y0), obj_r)115 116 mask = Image.new("L", (W, H), 0)117 obj_alpha = obj_r.split()[-1]118 mask.paste(obj_alpha, (x0, y0), obj_alpha)119 120 # Geometry proxy: the object RGB on a black canvas at its placed location.121 geometry_full = Image.new("RGB", (W, H), (0, 0, 0))122 geometry_full.paste(obj_r, (x0, y0), obj_r)123 124 return placed_rgb, mask, geometry_full125 126 127def get_mask_bbox(mask_pil, threshold=20):128 arr = np.array(mask_pil)129 ys, xs = np.where(arr > threshold)130 if ys.size == 0:131 return None132 return (xs.min(), ys.min(), xs.max() + 1, ys.max() + 1)133 134 135def get_smart_crop_bbox(mask_pil, min_ratio=0.02, max_ratio=0.3):136 bbox = get_mask_bbox(mask_pil)137 if bbox is None:138 s = MODEL_INPUT_RESOLUTION139 return (0, 0, s, s), s140 min_x, min_y, max_x, max_y = bbox141 mask_w, mask_h = max_x - min_x, max_y - min_y142 area = mask_w * mask_h143 side = int(math.sqrt(area / ((min_ratio + max_ratio) / 2.0)))144 side = max(side, max(mask_w, mask_h) + 40)145 cx = (min_x + max_x) // 2146 cy = (min_y + max_y) // 2147 half = side // 2148 return (cx - half, cy - half, cx - half + side, cy - half + side), side149 150 151def crop_and_pad(image, bbox, target_side):152 x1, y1, x2, y2 = bbox153 W, H = image.size154 valid = image.crop((max(0, x1), max(0, y1), min(W, x2), min(H, y2)))155 canvas = Image.new(image.mode, (target_side, target_side), 0)156 canvas.paste(valid, (max(0, -x1), max(0, -y1)))157 return canvas158 159 160def dilate_mask(mask_np, radius=10):161 import cv2162 163 m = (mask_np > 0).astype(np.uint8) * 255164 k = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (radius * 2 + 1, radius * 2 + 1))165 return (cv2.dilate(m, k, iterations=1) > 0).astype(np.uint8)166 167 168def refine_mask_holes(mask_bool, kernel_size=7):169 import cv2170 171 m = mask_bool.astype(np.uint8) * 255172 k = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (kernel_size, kernel_size))173 closed = cv2.morphologyEx(m, cv2.MORPH_CLOSE, k, iterations=2)174 contours, _ = cv2.findContours(closed, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)175 filled = np.zeros_like(closed)176 cv2.drawContours(filled, contours, -1, 255, thickness=cv2.FILLED)177 return filled > 127178 179 180def adain_color_fix(target_pil, source_pil, mask_pil):181 from torchvision.transforms import ToPILImage, ToTensor182 183 to_tensor = ToTensor()184 t = to_tensor(target_pil).unsqueeze(0)185 s = to_tensor(source_pil).unsqueeze(0)186 m = to_tensor(mask_pil).unsqueeze(0)187 eps = 1e-5188 res = t.clone()189 for ch in range(3):190 bg_idx = m[0, 0] < 0.1191 if bg_idx.sum() < 10:192 continue193 s_pix = s[0, ch][bg_idx]194 t_pix = t[0, ch][bg_idx]195 s_mean, s_std = s_pix.mean(), s_pix.std() + eps196 t_mean, t_std = t_pix.mean(), t_pix.std() + eps197 res[0, ch] = (t[0, ch] - t_mean) * (s_std / t_std) + s_mean198 return ToPILImage()(res.squeeze(0).clamp(0, 1))199 200 201def build_inputs(bg_pil, composite_full, mask_full, reference_ref, geometry_full):202 """Produce the model's 1024x1024 conditioning tensors from full-frame inputs."""203 target_res = MODEL_INPUT_RESOLUTION204 205 mask_np = np.array(mask_full)206 dilated01 = dilate_mask(mask_np, radius=10)207 dilated_pil = Image.fromarray(dilated01 * 255, mode="L")208 209 # Context image: full background with the (dilated) insertion region blacked.210 full_bg = np.array(bg_pil.convert("RGB"))211 context_image = Image.fromarray((full_bg * (1 - dilated01[:, :, None])).astype(np.uint8))212 213 ideal_bbox, target_side = get_smart_crop_bbox(dilated_pil)214 215 patch_composite = crop_and_pad(composite_full, ideal_bbox, target_side)216 patch_mask = crop_and_pad(dilated_pil, ideal_bbox, target_side)217 patch_geometry = crop_and_pad(geometry_full, ideal_bbox, target_side)218 patch_bg_ref = crop_and_pad(bg_pil.convert("RGB"), ideal_bbox, target_side)219 patch_mask_orig = crop_and_pad(Image.fromarray(mask_np), ideal_bbox, target_side)220 221 comp_arr = np.array(patch_composite)222 mask_dilated_arr = np.array(patch_mask) > 127223 mask_orig_arr = refine_mask_holes(np.array(patch_mask_orig) > 127, kernel_size=7)224 diff_region = mask_dilated_arr & (~mask_orig_arr)225 comp_arr[diff_region] = [0, 0, 0]226 patch_composite = Image.fromarray(comp_arr)227 228 composite_image = patch_composite.resize((target_res, target_res), Image.LANCZOS)229 model_input_mask = Image.fromarray(np.array(patch_mask).astype(np.uint8)).resize(230 (target_res, target_res), Image.NEAREST231 )232 geometry_image = patch_geometry.resize((target_res, target_res), Image.LANCZOS)233 background_reference_image = patch_bg_ref.resize((target_res, target_res), Image.LANCZOS)234 235 inpaint_mask = Image.fromarray(((np.array(model_input_mask) > 0) * 255).astype(np.uint8))236 237 return {238 "composite_image": composite_image,239 "inpaint_mask": inpaint_mask,240 "reference_image": reference_ref,241 "geometry_image": geometry_image,242 "context_image": context_image,243 "model_input_mask": model_input_mask,244 "background_reference_image": background_reference_image,245 "ideal_bbox": ideal_bbox,246 "target_side": target_side,247 }248 249 250def paste_back(bg_pil, generated_patch, inp):251 fixed = adain_color_fix(252 generated_patch, inp["background_reference_image"], inp["model_input_mask"]253 )254 fixed = fixed.resize((inp["target_side"], inp["target_side"]), Image.LANCZOS)255 x1, y1, x2, y2 = inp["ideal_bbox"]256 W, H = bg_pil.size257 pad_left = max(0, -x1)258 pad_top = max(0, -y1)259 valid_w = min(W, x2) - max(0, x1)260 valid_h = min(H, y2) - max(0, y1)261 patch_valid = fixed.crop((pad_left, pad_top, pad_left + valid_w, pad_top + valid_h))262 out = bg_pil.convert("RGB").copy()263 out.paste(patch_valid, (max(0, x1), max(0, y1)))264 return out265 266 267# ----------------------------------------------------------------------------268# Inference269# ----------------------------------------------------------------------------270 271def _estimate_duration(bg, obj, cx, cy, scale, seed, ref_scale, steps, *a, **k):272 # Measured ~12 s/step at 1024 when reference guidance is on (CFG doubles the273 # forward pass); ~half that when it is off. Plus fixed overhead for VAE /274 # rembg / cold worker init.275 try:276 steps = int(steps)277 except Exception:278 steps = 16279 try:280 ref_on = float(ref_scale) > 1.0281 except Exception:282 ref_on = True283 per_step = 12.5 if ref_on else 6.5284 return int(min(600, 45 + steps * per_step))285 286 287@spaces.GPU(duration=_estimate_duration)288def insert_object(289 bg: Image.Image,290 obj: Image.Image,291 cx: float,292 cy: float,293 scale: float,294 seed: int,295 ref_scale: float,296 steps: int,297 progress=gr.Progress(track_tqdm=True),298):299 """Insert a reference object into a background image with 3D-aware harmonization.300 301 Args:302 bg: Background scene image.303 obj: Reference object image (background is removed automatically).304 cx: Horizontal placement of the object center (0=left, 1=right).305 cy: Vertical placement of the object center (0=top, 1=bottom).306 scale: Object size as a fraction of the background's longest side.307 seed: Random seed for reproducibility.308 ref_scale: Reference guidance scale (identity preservation strength).309 steps: Number of inference steps.310 311 Returns:312 The composited image with the object inserted, and a preview of the raw313 2D placement used as geometric guidance.314 """315 if bg is None:316 raise gr.Error("Please provide a background image.")317 if obj is None:318 raise gr.Error("Please provide an object image.")319 320 t0 = time.perf_counter()321 bg = bg.convert("RGB")322 obj_rgba = segment_object(obj)323 324 reference_ref = center_reference(obj_rgba, out_size=MODEL_INPUT_RESOLUTION)325 placed_rgb, mask_full, geometry_full = place_object(bg, obj_rgba, cx, cy, scale)326 327 inp = build_inputs(bg, placed_rgb, mask_full, reference_ref, geometry_full)328 329 seed = int(seed)330 final_images = direct_pipeline(331 composite_image=inp["composite_image"],332 inpaint_mask=inp["inpaint_mask"],333 reference_image=inp["reference_image"],334 geometry_image=inp["geometry_image"],335 context_image=inp["context_image"],336 seed=seed,337 guidance_scale=30,338 num_inference_steps=int(steps),339 height=MODEL_INPUT_RESOLUTION,340 width=MODEL_INPUT_RESOLUTION,341 use_autocast=True,342 reference_guidance_scale=float(ref_scale),343 )344 generated_patch = final_images[0]345 result = paste_back(bg, generated_patch, inp)346 print(f"[insert_object] done in {time.perf_counter() - t0:.1f}s (steps={steps})")347 return result, placed_rgb348 349 350def randomize_seed():351 return random.randint(0, 2**31 - 1)352 353 354# ----------------------------------------------------------------------------355# UI356# ----------------------------------------------------------------------------357CSS = """358#col-container { max-width: 1200px; margin: 0 auto; }359.dark .gradio-container { color: var(--body-text-color); }360"""361 362INTRO = """363# DIRECT: 3D-Aware Object Insertion364 365Insert a reference **object** into a **background** scene with realistic,366harmonized results, powered by the [DIRECT](https://huggingface.co/superGong/DIRECT)367model (ICML 2026) — a FLUX.1-Fill-dev network guided by a decomposed visual proxy.368 369**How to use:** upload a background and an object image (its background is370removed automatically), choose *where* and *how big* to place it, then click **Insert**.371 372> **Note.** The full paper uses an interactive 3D viewer (TRELLIS + Viser) to pose a373> reconstructed 3D proxy of the object. That live 3D viewer cannot run inside a374> single-port Space, so this demo drives the same DIRECT model with a simpler375> **2D placement** (position + scale) as its geometric guidance.376 377[Paper](https://arxiv.org/abs/2606.06601) · [Project page](https://gong1130.github.io/DIRECT/) · [Code](https://github.com/Gong1130/DIRECT)378"""379 380with gr.Blocks(theme=gr.themes.Citrus(), css=CSS) as demo:381 with gr.Column(elem_id="col-container"):382 gr.Markdown(INTRO)383 with gr.Row():384 with gr.Column(scale=1):385 bg_input = gr.Image(label="Background image", type="pil", height=300)386 obj_input = gr.Image(label="Object image", type="pil", height=300)387 run_btn = gr.Button("Insert", variant="primary")388 with gr.Column(scale=1):389 out_result = gr.Image(label="Inserted result", type="pil", height=360)390 out_preview = gr.Image(label="2D placement (geometric guidance)", type="pil", height=240)391 392 with gr.Accordion("Placement & advanced settings", open=True):393 with gr.Row():394 cx = gr.Slider(0.0, 1.0, value=0.5, step=0.01, label="Horizontal position")395 cy = gr.Slider(0.0, 1.0, value=0.6, step=0.01, label="Vertical position")396 scale = gr.Slider(0.05, 0.9, value=0.35, step=0.01, label="Object size")397 with gr.Row():398 ref_scale = gr.Slider(1.0, 5.0, value=2.0, step=0.1, label="Reference guidance scale")399 steps = gr.Slider(12, 28, value=16, step=1, label="Inference steps")400 seed = gr.Number(label="Seed", value=42, precision=0)401 rand_btn = gr.Button("🎲 Randomize seed")402 403 gr.Examples(404 examples=[405 ["examples/bg_landscape.jpg", "examples/obj_ducks.jpg", 0.55, 0.70, 0.28, 42, 2.0, 16],406 ["examples/bg_tent.jpg", "examples/obj_dog.jpg", 0.45, 0.68, 0.30, 7, 2.0, 16],407 ["examples/bg_beach.jpg", "examples/obj_cake.jpg", 0.50, 0.72, 0.22, 123, 2.5, 16],408 ],409 inputs=[bg_input, obj_input, cx, cy, scale, seed, ref_scale, steps],410 outputs=[out_result, out_preview],411 fn=insert_object,412 cache_examples=True,413 cache_mode="lazy",414 )415 416 rand_btn.click(fn=randomize_seed, outputs=seed)417 run_btn.click(418 fn=insert_object,419 inputs=[bg_input, obj_input, cx, cy, scale, seed, ref_scale, steps],420 outputs=[out_result, out_preview],421 api_name="insert",422 )423 424if __name__ == "__main__":425 demo.launch(mcp_server=True)426