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)