Update watermark/engine.py

watermark/engine.py

@@ -1,226 +1,1125 @@
 from pathlib import Path
 import os
 import urllib.request
+
 import numpy as np
 import torch
 from PIL import Image, ImageDraw, ImageFilter
 
+
+# ============================================================
+# Global model cache
+# ============================================================
+
 MODEL_CACHE = None
+
 FLORENCE_MODEL = None
 FLORENCE_PROCESSOR = None
+
 LAMA_MODEL = None
 
+
+# ============================================================
+# Device / dtype
+# ============================================================
+
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+
+if DEVICE == "cuda":
+    DTYPE = torch.bfloat16
+else:
+    DTYPE = torch.float32
+
+
+# ============================================================
+# Writable cache directory
+# ============================================================
+
 def _get_writable_cache():
+    """
+    Find a writable directory.
+
+    HF Spaces may not allow writing to /data, so never assume
+    /data exists or is writable.
+    """
+
     candidates = []
-    if os.getenv("WATERMARK_CACHE_DIR"):
-        candidates.append(Path(os.environ["WATERMARK_CACHE_DIR"]))
-    candidates += [
-        Path("/data/.cache/watermark"),
-        Path("/tmp/watermark"),
-        Path.home() / ".cache" / "watermark",
-    ]
-    for p in candidates:
+
+    env_cache = os.getenv("WATERMARK_CACHE_DIR")
+
+    if env_cache:
+        candidates.append(Path(env_cache))
+
+    candidates.extend(
+        [
+            Path("/tmp/watermark"),
+            Path.home() / ".cache" / "watermark",
+        ]
+    )
+
+    for path in candidates:
+
         try:
-            p.mkdir(parents=True, exist_ok=True)
-            test = p / ".write_test"
-            test.write_text("ok", encoding="utf-8")
-            test.unlink(missing_ok=True)
-            print(f"Watermark cache: {p}")
-            return p
-        except (PermissionError, OSError):
-            pass
-    raise RuntimeError("No writable watermark cache directory available.")
+            path.mkdir(
+                parents=True,
+                exist_ok=True,
+            )
+
+            test_file = path / ".write_test"
+
+            test_file.write_text(
+                "ok",
+                encoding="utf-8",
+            )
+
+            try:
+                test_file.unlink()
+            except OSError:
+                pass
+
+            print(f"[Watermark] Cache directory: {path}")
+
+            return path
+
+        except (PermissionError, OSError) as e:
+
+            print(
+                f"[Watermark] Cache unavailable: "
+                f"{path} ({e})"
+            )
+
+    raise RuntimeError(
+        "No writable watermark cache directory available."
+    )
+
 
 def _cache_dir():
     global MODEL_CACHE
+
     if MODEL_CACHE is None:
         MODEL_CACHE = _get_writable_cache()
+
     return MODEL_CACHE
 
+
+# ============================================================
+# Florence-2 compatibility
+# ============================================================
+
 def _patch_florence_config():
-    # Florence-2 custom configuration from older revisions accesses
-    # self.forced_bos_token_id directly. Newer Transformers may not
-    # materialize that attribute on PretrainedConfig subclasses.
+    """
+    Compatibility patch for Florence-2 custom code.
+
+    Florence-2 was originally written against older versions
+    of Transformers.
+
+    Newer Transformers may access forced_bos_token_id before
+    the attribute exists.
+
+    This patch is harmless when the attribute already exists.
+    """
+
     try:
+
         from transformers import PretrainedConfig
-        if not hasattr(PretrainedConfig, "forced_bos_token_id"):
+
+        if not hasattr(
+            PretrainedConfig,
+            "forced_bos_token_id",
+        ):
+
             PretrainedConfig.forced_bos_token_id = None
-            print("Applied Florence-2 forced_bos_token_id compatibility patch.")
+
+            print(
+                "[Florence] Applied "
+                "forced_bos_token_id compatibility patch."
+            )
+
+    except Exception as e:
+
+        print(
+            f"[Florence] Config patch skipped: {e}"
+        )
+
+
+def _patch_florence_model_class():
+    """
+    Fix Florence-2 _supports_sdpa compatibility.
+
+    Some Florence-2 remote-code revisions define
+    _supports_sdpa as a property that accesses
+    self.language_model.
+
+    Transformers may query _supports_sdpa during
+    Florence2ForConditionalGeneration.__init__(),
+    before language_model has been created.
+
+    Therefore we force class-level attributes.
+    """
+
+    try:
+
+        from transformers.models.auto.modeling_auto import (
+            AutoModelForCausalLM,
+        )
+
+        # Importing the model through the dynamic module system
+        # happens only after from_pretrained(), so this patch is
+        # additionally applied after model class creation below.
+
+        return AutoModelForCausalLM
+
     except Exception as e:
-        print(f"Florence compatibility patch skipped: {e}")
+
+        print(
+            f"[Florence] Model compatibility setup skipped: {e}"
+        )
+
+        return None
+
+
+# ============================================================
+# Load Florence + LaMa
+# ============================================================
 
 def _load_models():
-    global FLORENCE_MODEL, FLORENCE_PROCESSOR, LAMA_MODEL
-    if FLORENCE_MODEL is not None and LAMA_MODEL is not None:
+
+    global FLORENCE_MODEL
+    global FLORENCE_PROCESSOR
+    global LAMA_MODEL
+
+    # Already loaded
+    if (
+        FLORENCE_MODEL is not None
+        and FLORENCE_PROCESSOR is not None
+        and LAMA_MODEL is not None
+    ):
         return
 
+    print(
+        f"[Watermark] Device: {DEVICE}"
+    )
+
+    print(
+        f"[Watermark] Dtype: {DTYPE}"
+    )
+
+    # --------------------------------------------------------
+    # Florence compatibility
+    # --------------------------------------------------------
+
     _patch_florence_config()
 
-    from transformers import AutoProcessor, AutoModelForCausalLM
+    # --------------------------------------------------------
+    # Transformers
+    # --------------------------------------------------------
+
+    from transformers import (
+        AutoProcessor,
+        AutoModelForCausalLM,
+    )
 
     cache = _cache_dir()
+
+    hf_cache = cache / "hf"
+
+    hf_cache.mkdir(
+        parents=True,
+        exist_ok=True,
+    )
+
     model_id = os.getenv(
         "FLORENCE_MODEL_ID",
         "microsoft/Florence-2-large",
     )
 
-    print(f"Loading Florence-2: {model_id}")
+    print(
+        f"[Florence] Loading model: {model_id}"
+    )
+
+    # --------------------------------------------------------
+    # Processor
+    # --------------------------------------------------------
 
-    FLORENCE_PROCESSOR = AutoProcessor.from_pretrained(
-        model_id,
-        cache_dir=str(cache / "hf"),
-        trust_remote_code=True,
+    FLORENCE_PROCESSOR = (
+        AutoProcessor.from_pretrained(
+            model_id,
+            cache_dir=str(hf_cache),
+            trust_remote_code=True,
+        )
     )
 
-    FLORENCE_MODEL = AutoModelForCausalLM.from_pretrained(
-        model_id,
-        torch_dtype=torch.bfloat16,
-        cache_dir=str(cache / "hf"),
-        trust_remote_code=True,
-    ).to("cuda").eval()
+    print(
+        "[Florence] Processor loaded."
+    )
+
+    # --------------------------------------------------------
+    # Florence model
+    # --------------------------------------------------------
+
+    model_kwargs = {
+        "cache_dir": str(hf_cache),
+        "trust_remote_code": True,
+    }
+
+    if DEVICE == "cuda":
+        model_kwargs["torch_dtype"] = DTYPE
+    else:
+        model_kwargs["torch_dtype"] = torch.float32
+
+    # IMPORTANT:
+    #
+    # eager avoids the SDPA compatibility path.
+    #
+    # This is particularly useful if the environment contains
+    # a newer Transformers version or newer Torch.
+    #
+    model_kwargs["attn_implementation"] = "eager"
+
+    print(
+        "[Florence] Loading model weights..."
+    )
+
+    try:
+
+        FLORENCE_MODEL = (
+            AutoModelForCausalLM.from_pretrained(
+                model_id,
+                **model_kwargs,
+            )
+        )
+
+    except TypeError as e:
+
+        # Some older Transformers versions may not accept
+        # attn_implementation.
+
+        if "attn_implementation" in str(e):
+
+            print(
+                "[Florence] "
+                "attn_implementation unsupported; "
+                "retrying without it."
+            )
+
+            model_kwargs.pop(
+                "attn_implementation",
+                None,
+            )
+
+            FLORENCE_MODEL = (
+                AutoModelForCausalLM.from_pretrained(
+                    model_id,
+                    **model_kwargs,
+                )
+            )
+
+        else:
+            raise
+
+    # --------------------------------------------------------
+    # Explicit Florence _supports_sdpa compatibility
+    # --------------------------------------------------------
+
+    try:
+
+        model_cls = FLORENCE_MODEL.__class__
+
+        # Force class-level values.
+
+        if (
+            not hasattr(model_cls, "_supports_sdpa")
+            or isinstance(
+                getattr(
+                    model_cls,
+                    "_supports_sdpa",
+                    None,
+                ),
+                property,
+            )
+        ):
+
+            try:
+                model_cls._supports_sdpa = True
+            except Exception:
+                pass
+
+        if (
+            not hasattr(
+                model_cls,
+                "_supports_flash_attn_2",
+            )
+            or isinstance(
+                getattr(
+                    model_cls,
+                    "_supports_flash_attn_2",
+                    None,
+                ),
+                property,
+            )
+        ):
+
+            try:
+                model_cls._supports_flash_attn_2 = False
+            except Exception:
+                pass
+
+    except Exception as e:
+
+        print(
+            f"[Florence] "
+            f"Class compatibility patch skipped: {e}"
+        )
+
+    # --------------------------------------------------------
+    # Move to GPU
+    # --------------------------------------------------------
+
+    FLORENCE_MODEL = (
+        FLORENCE_MODEL
+        .to(DEVICE)
+        .eval()
+    )
+
+    print(
+        "[Florence] Model loaded."
+    )
+
+    # --------------------------------------------------------
+    # LaMa
+    # --------------------------------------------------------
 
     lama_path = cache / "big-lama.pt"
+
     if not lama_path.exists():
-        url = os.getenv(
+
+        lama_url = os.getenv(
             "LAMA_URL",
             "https://github.com/Sanster/models/releases/download/add_big_lama/big-lama.pt",
         )
-        print("Downloading LaMa model...")
-        urllib.request.urlretrieve(url, lama_path)
+
+        print(
+            "[LaMa] Downloading model..."
+        )
+
+        print(
+            f"[LaMa] URL: {lama_url}"
+        )
+
+        urllib.request.urlretrieve(
+            lama_url,
+            lama_path,
+        )
+
+        print(
+            "[LaMa] Download complete."
+        )
+
+    print(
+        "[LaMa] Loading model..."
+    )
 
     LAMA_MODEL = torch.jit.load(
         str(lama_path),
-        map_location="cuda",
+        map_location=DEVICE,
     ).eval()
 
-    print("Watermark models loaded.")
+    print(
+        "[LaMa] Model loaded."
+    )
+
+    print(
+        "[Watermark] All watermark models loaded."
+    )
+
+
+# ============================================================
+# Florence watermark detection
+# ============================================================
 
 def _boxes(image, prompt):
+
     _load_models()
 
+    if image is None:
+        return []
+
+    image = image.convert("RGB")
+
+    prompt = (
+        str(prompt).strip()
+        if prompt
+        else "watermark"
+    )
+
+    # Florence-2 open vocabulary detection
     task = "<OPEN_VOCABULARY_DETECTION>"
+
     inputs = FLORENCE_PROCESSOR(
         text=task + prompt,
         images=image,
         return_tensors="pt",
     )
-    inputs = {
-        k: (v.to("cuda") if hasattr(v, "to") else v)
-        for k, v in inputs.items()
-    }
+
+    # Move tensors to device
+    prepared_inputs = {}
+
+    for key, value in inputs.items():
+
+        if hasattr(value, "to"):
+
+            if (
+                DEVICE == "cuda"
+                and torch.is_floating_point(value)
+            ):
+                value = value.to(
+                    DEVICE,
+                    dtype=DTYPE,
+                )
+            else:
+                value = value.to(DEVICE)
+
+        prepared_inputs[key] = value
+
+    # --------------------------------------------------------
+    # Florence generation
+    # --------------------------------------------------------
 
     with torch.inference_mode():
+
         generated = FLORENCE_MODEL.generate(
-            input_ids=inputs.get("input_ids"),
-            pixel_values=inputs.get("pixel_values"),
-            max_new_tokens=1024,
-            num_beams=3,
+
+            input_ids=prepared_inputs.get(
+                "input_ids"
+            ),
+
+            pixel_values=prepared_inputs.get(
+                "pixel_values"
+            ),
+
+            max_new_tokens=512,
+
+            # Beam search is unnecessary for this task.
+            # Reducing to 1 significantly lowers ZeroGPU cost.
+            num_beams=1,
+
+            do_sample=False,
         )
 
-    decoded = FLORENCE_PROCESSOR.batch_decode(
-        generated,
-        skip_special_tokens=False,
-    )[0]
+    # --------------------------------------------------------
+    # Decode
+    # --------------------------------------------------------
+
+    decoded = (
+        FLORENCE_PROCESSOR.batch_decode(
+            generated,
+            skip_special_tokens=False,
+        )[0]
+    )
+
+    # --------------------------------------------------------
+    # Parse Florence result
+    # --------------------------------------------------------
+
+    parsed = (
+        FLORENCE_PROCESSOR.post_process_generation(
+            decoded,
+            task=task,
+            image_size=image.size,
+        )
+    )
 
-    parsed = FLORENCE_PROCESSOR.post_process_generation(
-        decoded,
-        task=task,
-        image_size=image.size,
+    data = parsed.get(
+        task,
+        parsed,
     )
-    data = parsed.get(task, parsed)
 
     if not isinstance(data, dict):
         return []
 
-    return data.get("bboxes", []) or []
+    boxes = data.get(
+        "bboxes",
+        [],
+    )
+
+    if not boxes:
+        return []
+
+    # --------------------------------------------------------
+    # Validate boxes
+    # --------------------------------------------------------
+
+    valid_boxes = []
+
+    for box in boxes:
+
+        if not isinstance(
+            box,
+            (list, tuple),
+        ):
+            continue
+
+        if len(box) != 4:
+            continue
+
+        try:
+
+            x1, y1, x2, y2 = map(
+                float,
+                box,
+            )
+
+        except Exception:
+            continue
+
+        if x2 <= x1 or y2 <= y1:
+            continue
+
+        valid_boxes.append(
+            [
+                x1,
+                y1,
+                x2,
+                y2,
+            ]
+        )
+
+    return valid_boxes
+
+
+# ============================================================
+# Create watermark mask
+# ============================================================
+
+def _mask(
+    image,
+    boxes,
+    max_pct,
+):
+
+    image = image.convert("RGB")
+
+    mask = Image.new(
+        "L",
+        image.size,
+        0,
+    )
 
-def _mask(image, boxes, max_pct):
-    mask = Image.new("L", image.size, 0)
     draw = ImageDraw.Draw(mask)
-    total = image.width * image.height
+
+    total_area = (
+        image.width
+        * image.height
+    )
+
+    try:
+        max_pct = float(max_pct)
+    except Exception:
+        max_pct = 30.0
+
+    max_pct = max(
+        0.1,
+        min(
+            max_pct,
+            100.0,
+        ),
+    )
 
     for box in boxes:
+
         if len(box) != 4:
             continue
-        x1, y1, x2, y2 = map(float, box)
-        x1 = max(0, min(image.width, x1))
-        x2 = max(0, min(image.width, x2))
-        y1 = max(0, min(image.height, y1))
-        y2 = max(0, min(image.height, y2))
-        area = max(0, x2 - x1) * max(0, y2 - y1)
 
-        if total and area / total * 100 <= max_pct:
-            draw.rectangle((x1, y1, x2, y2), fill=255)
+        try:
+
+            x1, y1, x2, y2 = map(
+                float,
+                box,
+            )
+
+        except Exception:
+            continue
+
+        x1 = max(
+            0,
+            min(
+                image.width,
+                x1,
+            ),
+        )
+
+        x2 = max(
+            0,
+            min(
+                image.width,
+                x2,
+            ),
+        )
+
+        y1 = max(
+            0,
+            min(
+                image.height,
+                y1,
+            ),
+        )
+
+        y2 = max(
+            0,
+            min(
+                image.height,
+                y2,
+            ),
+        )
+
+        area = (
+            max(
+                0,
+                x2 - x1,
+            )
+            *
+            max(
+                0,
+                y2 - y1,
+            )
+        )
+
+        if (
+            total_area > 0
+            and area / total_area * 100
+            <= max_pct
+        ):
+
+            draw.rectangle(
+                (
+                    x1,
+                    y1,
+                    x2,
+                    y2,
+                ),
+                fill=255,
+            )
+
+    # Slight expansion to ensure watermark edges
+    # are also included.
+    try:
+
+        mask = mask.filter(
+            ImageFilter.MaxFilter(7)
+        )
+
+    except Exception:
+        pass
+
+    return mask
+
 
-    return mask.filter(ImageFilter.MaxFilter(7))
+# ============================================================
+# LaMa inpainting
+# ============================================================
+
+def _lama(
+    image,
+    mask,
+    boxes=None,
+):
+    """
+    Run LaMa inpainting.
+
+    If bounding boxes are available, process only local
+    regions. This is substantially cheaper for video.
+    """
 
-def _lama(image, mask, boxes=None):
-    """Run LaMa on watermark regions rather than the whole frame when boxes exist."""
     _load_models()
 
     image = image.convert("RGB")
     mask = mask.convert("L")
 
-    # Region-based inpainting is much cheaper for video than sending a full
-    # 1080p/4K frame through LaMa. Use a modest padding around each box.
+    # --------------------------------------------------------
+    # Region-based processing
+    # --------------------------------------------------------
+
     if boxes:
+
         result = image.copy()
-        pad = max(16, int(min(image.size) * 0.02))
+
+        # Padding around watermark region.
+        pad = max(
+            16,
+            int(
+                min(image.size)
+                * 0.02
+            ),
+        )
+
         for box in boxes:
+
             if len(box) != 4:
                 continue
-            x1, y1, x2, y2 = map(float, box)
-            left = max(0, int(x1) - pad)
-            top = max(0, int(y1) - pad)
-            right = min(image.width, int(x2) + pad)
-            bottom = min(image.height, int(y2) + pad)
-            if right <= left or bottom <= top:
+
+            try:
+
+                x1, y1, x2, y2 = map(
+                    float,
+                    box,
+                )
+
+            except Exception:
                 continue
-            crop = image.crop((left, top, right, bottom))
-            crop_mask = mask.crop((left, top, right, bottom))
-            crop_np = np.asarray(crop, dtype=np.float32) / 255.0
-            mask_np = np.asarray(crop_mask, dtype=np.float32) / 255.0
-            image_t = torch.from_numpy(crop_np).permute(2,0,1).unsqueeze(0).cuda() * 2.0 - 1.0
-            mask_t = torch.from_numpy(mask_np).unsqueeze(0).unsqueeze(0).cuda()
-            mask_t = (mask_t > 0.5).float()
+
+            left = max(
+                0,
+                int(x1) - pad,
+            )
+
+            top = max(
+                0,
+                int(y1) - pad,
+            )
+
+            right = min(
+                image.width,
+                int(x2) + pad,
+            )
+
+            bottom = min(
+                image.height,
+                int(y2) + pad,
+            )
+
+            if (
+                right <= left
+                or bottom <= top
+            ):
+                continue
+
+            crop = image.crop(
+                (
+                    left,
+                    top,
+                    right,
+                    bottom,
+                )
+            )
+
+            crop_mask = mask.crop(
+                (
+                    left,
+                    top,
+                    right,
+                    bottom,
+                )
+            )
+
+            # ------------------------------------------------
+            # Convert image
+            # ------------------------------------------------
+
+            crop_np = (
+                np.asarray(
+                    crop,
+                    dtype=np.float32,
+                )
+                / 255.0
+            )
+
+            mask_np = (
+                np.asarray(
+                    crop_mask,
+                    dtype=np.float32,
+                )
+                / 255.0
+            )
+
+            image_t = (
+                torch.from_numpy(
+                    crop_np
+                )
+                .permute(
+                    2,
+                    0,
+                    1,
+                )
+                .unsqueeze(0)
+                .to(DEVICE)
+                * 2.0
+                - 1.0
+            )
+
+            mask_t = (
+                torch.from_numpy(
+                    mask_np
+                )
+                .unsqueeze(0)
+                .unsqueeze(0)
+                .to(DEVICE)
+            )
+
+            mask_t = (
+                mask_t > 0.5
+            ).float()
+
+            # ------------------------------------------------
+            # LaMa inference
+            # ------------------------------------------------
+
             with torch.inference_mode():
-                out = LAMA_MODEL(image_t, mask_t)
-            if isinstance(out, (tuple, list)):
-                out = out[0]
-            out = ((out.clamp(-1, 1) + 1.0) / 2.0).squeeze(0).permute(1,2,0).float().cpu().numpy()
-            restored = Image.fromarray((out * 255).clip(0,255).astype(np.uint8))
-            result.paste(restored, (left, top))
+
+                output = LAMA_MODEL(
+                    image_t,
+                    mask_t,
+                )
+
+            if isinstance(
+                output,
+                (tuple, list),
+            ):
+                output = output[0]
+
+            # ------------------------------------------------
+            # Convert back to PIL
+            # ------------------------------------------------
+
+            output = (
+                (
+                    output.clamp(
+                        -1,
+                        1,
+                    )
+                    + 1.0
+                )
+                / 2.0
+            )
+
+            output = (
+                output
+                .squeeze(0)
+                .permute(
+                    1,
+                    2,
+                    0,
+                )
+                .float()
+                .cpu()
+                .numpy()
+            )
+
+            restored = Image.fromarray(
+                (
+                    output * 255
+                )
+                .clip(
+                    0,
+                    255,
+                )
+                .astype(
+                    np.uint8
+                )
+            )
+
+            result.paste(
+                restored,
+                (
+                    left,
+                    top,
+                ),
+            )
+
         return result
 
-    image_np = np.asarray(image, dtype=np.float32) / 255.0
-    mask_np = np.asarray(mask, dtype=np.float32) / 255.0
-    image_t = torch.from_numpy(image_np).permute(2,0,1).unsqueeze(0).cuda() * 2.0 - 1.0
-    mask_t = torch.from_numpy(mask_np).unsqueeze(0).unsqueeze(0).cuda()
-    mask_t = (mask_t > 0.5).float()
+    # ========================================================
+    # Full-frame fallback
+    # ========================================================
+
+    image_np = (
+        np.asarray(
+            image,
+            dtype=np.float32,
+        )
+        / 255.0
+    )
+
+    mask_np = (
+        np.asarray(
+            mask,
+            dtype=np.float32,
+        )
+        / 255.0
+    )
+
+    image_t = (
+        torch.from_numpy(
+            image_np
+        )
+        .permute(
+            2,
+            0,
+            1,
+        )
+        .unsqueeze(0)
+        .to(DEVICE)
+        * 2.0
+        - 1.0
+    )
+
+    mask_t = (
+        torch.from_numpy(
+            mask_np
+        )
+        .unsqueeze(0)
+        .unsqueeze(0)
+        .to(DEVICE)
+    )
+
+    mask_t = (
+        mask_t > 0.5
+    ).float()
+
     with torch.inference_mode():
-        result = LAMA_MODEL(image_t, mask_t)
-    if isinstance(result, (tuple, list)):
-        result = result[0]
-    result = ((result.clamp(-1,1)+1.0)/2.0).squeeze(0).permute(1,2,0).float().cpu().numpy()
-    return Image.fromarray((result*255).clip(0,255).astype(np.uint8))
 
-def preview(image, max_bbox_percent=30, detection_prompt="watermark"):
+        output = LAMA_MODEL(
+            image_t,
+            mask_t,
+        )
+
+    if isinstance(
+        output,
+        (tuple, list),
+    ):
+        output = output[0]
+
+    output = (
+        (
+            output.clamp(
+                -1,
+                1,
+            )
+            + 1.0
+        )
+        / 2.0
+    )
+
+    output = (
+        output
+        .squeeze(0)
+        .permute(
+            1,
+            2,
+            0,
+        )
+        .float()
+        .cpu()
+        .numpy()
+    )
+
+    return Image.fromarray(
+        (
+            output * 255
+        )
+        .clip(
+            0,
+            255,
+        )
+        .astype(
+            np.uint8
+        )
+    )
+
+
+# ============================================================
+# Preview
+# ============================================================
+
+def preview(
+    image,
+    max_bbox_percent=30,
+    detection_prompt="watermark",
+):
+
     if image is None:
-        raise ValueError("Please upload an image.")
+        raise ValueError(
+            "Please upload an image."
+        )
+
+    image = image.convert("RGB")
+
+    boxes = _boxes(
+        image,
+        detection_prompt,
+    )
+
+    # Apply area filtering for preview as well.
+    filtered_boxes = []
 
-    boxes = _boxes(image, detection_prompt)
-    out = image.convert("RGB").copy()
-    draw = ImageDraw.Draw(out)
+    total_area = (
+        image.width
+        * image.height
+    )
+
+    try:
+        max_pct = float(
+            max_bbox_percent
+        )
+    except Exception:
+        max_pct = 30.0
 
     for box in boxes:
-        if len(box) == 4:
-            draw.rectangle(
-                tuple(map(float, box)),
-                outline="red",
-                width=4,
+
+        if len(box) != 4:
+            continue
+
+        x1, y1, x2, y2 = map(
+            float,
+            box,
+        )
+
+        area = (
+            max(
+                0,
+                x2 - x1,
+            )
+            *
+            max(
+                0,
+                y2 - y1,
             )
+        )
+
+        if (
+            total_area > 0
+            and area / total_area * 100
+            <= max_pct
+        ):
+
+            filtered_boxes.append(
+                box
+            )
+
+    output = image.copy()
+
+    draw = ImageDraw.Draw(
+        output
+    )
+
+    for box in filtered_boxes:
+
+        draw.rectangle(
+            tuple(
+                map(
+                    float,
+                    box,
+                )
+            ),
+            outline="red",
+            width=4,
+        )
+
+    return output
 
-    return out
+
+# ============================================================
+# Image watermark removal
+# ============================================================
 
 def remove_image(
     image,
@@ -228,19 +1127,54 @@ def remove_image(
     detection_prompt="watermark",
     transparent=False,
 ):
+
     if image is None:
-        raise ValueError("Please upload an image.")
+        raise ValueError(
+            "Please upload an image."
+        )
 
-    boxes = _boxes(image, detection_prompt)
-    mask = _mask(image, boxes, max_bbox_percent)
+    image = image.convert("RGB")
+
+    boxes = _boxes(
+        image,
+        detection_prompt,
+    )
 
+    mask = _mask(
+        image,
+        boxes,
+        max_bbox_percent,
+    )
+
+    # Nothing detected
     if not boxes:
-        return image.convert("RGBA" if transparent else "RGB"), mask
 
-    result = _lama(image, mask)
+        output = image.convert(
+            "RGBA"
+            if transparent
+            else "RGB"
+        )
+
+        return output, mask
+
+    # IMPORTANT:
+    # Pass boxes so LaMa performs local inpainting.
+    result = _lama(
+        image,
+        mask,
+        boxes,
+    )
 
     if transparent:
-        result = result.convert("RGBA")
-        result.putalpha(mask.point(lambda p: 255 - p))
+
+        result = result.convert(
+            "RGBA"
+        )
+
+        result.putalpha(
+            mask.point(
+                lambda p: 255 - p
+            )
+        )
 
     return result, mask