Update watermark/engine.py

watermark/engine.py

@@ -1,16 +1,14 @@
 from pathlib import Path
 import os
 import urllib.request
-import types
 
 import numpy as np
 import torch
-
 from PIL import Image, ImageDraw, ImageFilter
 
 
 # ============================================================
-# Global model state
+# Global model cache
 # ============================================================
 
 MODEL_CACHE = None
@@ -22,7 +20,37 @@ LAMA_MODEL = None
 
 
 # ============================================================
-# Cache
+# Device / dtype helpers
+# ============================================================
+
+def _get_device():
+    if torch.cuda.is_available():
+        return torch.device("cuda")
+    return torch.device("cpu")
+
+
+def _get_model_dtype():
+    """
+    Select a dtype that is actually usable by the current GPU.
+
+    ZeroGPU normally provides CUDA.
+    Prefer BF16 when supported, otherwise FP16.
+    """
+
+    if not torch.cuda.is_available():
+        return torch.float32
+
+    try:
+        if torch.cuda.is_bf16_supported():
+            return torch.bfloat16
+    except Exception:
+        pass
+
+    return torch.float16
+
+
+# ============================================================
+# Writable cache
 # ============================================================
 
 def _get_writable_cache():
@@ -53,9 +81,10 @@ def _get_writable_cache():
                 encoding="utf-8",
             )
 
-            test.unlink(
-                missing_ok=True
-            )
+            try:
+                test.unlink()
+            except Exception:
+                pass
 
             print(
                 f"[Watermark] Cache directory: {p}"
@@ -63,10 +92,7 @@ def _get_writable_cache():
 
             return p
 
-        except (
-            PermissionError,
-            OSError,
-        ):
+        except (PermissionError, OSError):
             pass
 
     raise RuntimeError(
@@ -84,20 +110,21 @@ def _cache_dir():
 
 
 # ============================================================
-# Florence-2 compatibility
+# Florence compatibility patches
 # ============================================================
 
-def _patch_florence_config():
+def _patch_florence_compatibility():
     """
-    Compatibility patch for older Florence-2 custom configuration.
-
-    Some Florence-2 revisions access:
-        self.forced_bos_token_id
+    Compatibility layer for Microsoft Florence-2 custom code.
 
-    Newer Transformers configurations may not expose this
-    attribute directly.
+    Handles older Florence-2 remote-code expectations when used
+    with newer Transformers versions.
     """
 
+    # --------------------------------------------------------
+    # forced_bos_token_id
+    # --------------------------------------------------------
+
     try:
         from transformers import PretrainedConfig
 
@@ -108,239 +135,46 @@ def _patch_florence_config():
             PretrainedConfig.forced_bos_token_id = None
 
             print(
-                "[Watermark] Applied "
-                "forced_bos_token_id compatibility patch."
+                "[Florence] "
+                "Applied forced_bos_token_id compatibility patch."
             )
 
     except Exception as e:
         print(
-            "[Watermark] Florence config patch skipped:",
-            e,
+            "[Florence] "
+            f"forced_bos_token_id patch skipped: {e}"
         )
 
-
-# ============================================================
-# Florence generation compatibility patch
-# ============================================================
-
-def _patch_florence_generation(model):
-    """
-    Patch Florence-2 language_model.prepare_inputs_for_generation.
-
-    Problem:
-
-        past_key_values[0][0].shape
-
-    can fail with:
-
-        AttributeError:
-        'NoneType' object has no attribute 'shape'
-
-    This happens because newer Transformers generation/cache
-    behavior can provide a cache structure whose first element
-    is None, while older Florence-2 remote code assumes a tensor.
-
-    We normalize malformed/empty caches to None and explicitly
-    handle the first generation step.
-    """
+    # --------------------------------------------------------
+    # _supports_sdpa
+    #
+    # Some Florence remote-code revisions do not define this
+    # attribute while newer Transformers expects it.
+    # --------------------------------------------------------
 
     try:
+        from transformers import PreTrainedModel
 
-        language_model = getattr(
-            model,
-            "language_model",
-            None,
-        )
-
-        if language_model is None:
-            print(
-                "[Watermark] No Florence language_model found; "
-                "generation patch skipped."
-            )
-            return
-
-        original_method = getattr(
-            language_model,
-            "prepare_inputs_for_generation",
-            None,
-        )
-
-        if original_method is None:
-            print(
-                "[Watermark] Florence language model does not "
-                "provide prepare_inputs_for_generation."
-            )
-            return
-
-        # Avoid installing the patch twice.
-        if getattr(
-            language_model,
-            "_watermark_generation_patched",
-            False,
-        ):
-            return
-
-        def patched_prepare_inputs(
-            self,
-            decoder_input_ids,
-            past_key_values=None,
-            attention_mask=None,
-            decoder_attention_mask=None,
-            head_mask=None,
-            decoder_head_mask=None,
-            cross_attn_head_mask=None,
-            use_cache=None,
-            encoder_outputs=None,
-            **kwargs,
+        if not hasattr(
+            PreTrainedModel,
+            "_supports_sdpa",
         ):
-            """
-            Robust Florence-2 generation input preparation.
-            """
-
-            # ------------------------------------------------
-            # Normalize malformed cache.
-            # ------------------------------------------------
-
-            cache_valid = False
-
-            if past_key_values is not None:
-
-                try:
-
-                    if len(past_key_values) > 0:
-
-                        first_layer = past_key_values[0]
-
-                        if first_layer is not None:
-
-                            if (
-                                isinstance(
-                                    first_layer,
-                                    (tuple, list),
-                                )
-                                and len(first_layer) >= 2
-                            ):
-
-                                key_tensor = first_layer[0]
-
-                                value_tensor = first_layer[1]
-
-                                if (
-                                    key_tensor is not None
-                                    and value_tensor is not None
-                                    and hasattr(
-                                        key_tensor,
-                                        "shape",
-                                    )
-                                    and hasattr(
-                                        value_tensor,
-                                        "shape",
-                                    )
-                                ):
-                                    cache_valid = True
-
-                except Exception:
-                    cache_valid = False
+            PreTrainedModel._supports_sdpa = False
 
-            # ------------------------------------------------
-            # Empty / malformed cache.
-            #
-            # Treat this as the first generation step.
-            # ------------------------------------------------
-
-            if not cache_valid:
-
-                past_key_values = None
-
-                remove_prefix_length = max(
-                    0,
-                    decoder_input_ids.shape[1] - 1,
-                )
-
-            else:
-
-                try:
-
-                    past_length = int(
-                        past_key_values[0][0].shape[2]
-                    )
-
-                except Exception:
-
-                    past_key_values = None
-
-                    remove_prefix_length = max(
-                        0,
-                        decoder_input_ids.shape[1] - 1,
-                    )
-
-                else:
-
-                    if (
-                        decoder_input_ids.shape[1]
-                        > past_length
-                    ):
-                        remove_prefix_length = (
-                            past_length
-                        )
-                    else:
-                        remove_prefix_length = (
-                            decoder_input_ids.shape[1]
-                            - 1
-                        )
-
-            # ------------------------------------------------
-            # Keep only the required decoder tokens.
-            # ------------------------------------------------
-
-            decoder_input_ids = (
-                decoder_input_ids[
-                    :,
-                    remove_prefix_length:
-                ]
+            print(
+                "[Florence] "
+                "Applied _supports_sdpa compatibility patch."
             )
 
-            # Never pass malformed cache.
-            if not cache_valid:
-                past_key_values = None
-
-            return {
-                "input_ids": None,
-                "encoder_outputs": encoder_outputs,
-                "past_key_values": past_key_values,
-                "decoder_input_ids": decoder_input_ids,
-                "attention_mask": attention_mask,
-                "decoder_attention_mask": decoder_attention_mask,
-                "head_mask": head_mask,
-                "decoder_head_mask": decoder_head_mask,
-                "cross_attn_head_mask": cross_attn_head_mask,
-                "use_cache": use_cache,
-            }
-
-        language_model.prepare_inputs_for_generation = (
-            types.MethodType(
-                patched_prepare_inputs,
-                language_model,
-            )
-        )
-
-        language_model._watermark_generation_patched = True
-
-        print(
-            "[Watermark] Florence generation compatibility "
-            "patch installed."
-        )
-
     except Exception as e:
-
         print(
-            "[Watermark] Generation patch failed:",
-            repr(e),
+            "[Florence] "
+            f"_supports_sdpa patch skipped: {e}"
         )
 
 
 # ============================================================
-# Florence model loading
+# Load models
 # ============================================================
 
 def _load_models():
@@ -349,10 +183,7 @@ def _load_models():
     global FLORENCE_PROCESSOR
     global LAMA_MODEL
 
-    # --------------------------------------------------------
     # Already loaded
-    # --------------------------------------------------------
-
     if (
         FLORENCE_MODEL is not None
         and FLORENCE_PROCESSOR is not None
@@ -360,11 +191,7 @@ def _load_models():
     ):
         return
 
-    # --------------------------------------------------------
-    # Florence compatibility
-    # --------------------------------------------------------
-
-    _patch_florence_config()
+    _patch_florence_compatibility()
 
     from transformers import (
         AutoProcessor,
@@ -378,13 +205,46 @@ def _load_models():
         "microsoft/Florence-2-large",
     )
 
+    device = _get_device()
+    dtype = _get_model_dtype()
+
     print(
-        f"[Watermark] Loading Florence-2: {model_id}"
+        "[Florence] =================================="
     )
 
-    # ========================================================
+    print(
+        f"[Florence] Model: {model_id}"
+    )
+
+    print(
+        f"[Florence] Device: {device}"
+    )
+
+    print(
+        f"[Florence] CUDA available: "
+        f"{torch.cuda.is_available()}"
+    )
+
+    if torch.cuda.is_available():
+        try:
+            print(
+                f"[Florence] BF16 supported: "
+                f"{torch.cuda.is_bf16_supported()}"
+            )
+        except Exception:
+            pass
+
+    print(
+        f"[Florence] dtype: {dtype}"
+    )
+
+    print(
+        "[Florence] =================================="
+    )
+
+    # --------------------------------------------------------
     # Processor
-    # ========================================================
+    # --------------------------------------------------------
 
     FLORENCE_PROCESSOR = (
         AutoProcessor.from_pretrained(
@@ -396,134 +256,69 @@ def _load_models():
         )
     )
 
-    print(
-        "[Watermark] Florence processor loaded."
-    )
-
-    # ========================================================
-    # Model
-    # ========================================================
+    # --------------------------------------------------------
+    # Florence model
+    # --------------------------------------------------------
 
     model_kwargs = {
-        "torch_dtype": torch.bfloat16,
         "cache_dir": str(
             cache / "hf"
         ),
         "trust_remote_code": True,
     }
 
-    # --------------------------------------------------------
-    # Prefer eager attention.
-    #
-    # This avoids the earlier:
-    #
-    # AttributeError:
-    # Florence2ForConditionalGeneration
-    # object has no attribute '_supports_sdpa'
-    # --------------------------------------------------------
-
-    try:
+    # Newer Transformers supports torch_dtype.
+    model_kwargs["torch_dtype"] = dtype
 
-        FLORENCE_MODEL = (
-            AutoModelForCausalLM.from_pretrained(
-                model_id,
-                attn_implementation="eager",
-                **model_kwargs,
-            )
-        )
-
-        print(
-            "[Watermark] Florence loaded with "
-            "eager attention."
-        )
-
-    except TypeError:
-
-        # Older Transformers may not accept
-        # attn_implementation.
-
-        print(
-            "[Watermark] attn_implementation is not "
-            "supported; loading Florence without it."
-        )
-
-        FLORENCE_MODEL = (
-            AutoModelForCausalLM.from_pretrained(
-                model_id,
-                **model_kwargs,
-            )
+    FLORENCE_MODEL = (
+        AutoModelForCausalLM.from_pretrained(
+            model_id,
+            **model_kwargs,
         )
+    )
 
-    # --------------------------------------------------------
-    # CUDA
-    # --------------------------------------------------------
-
+    # Explicit device placement
     FLORENCE_MODEL = (
         FLORENCE_MODEL
-        .to("cuda")
+        .to(device)
         .eval()
     )
 
-    # ========================================================
-    # Generation configuration
-    # ========================================================
+    # --------------------------------------------------------
+    # Extra compatibility patch directly on the model
+    # --------------------------------------------------------
 
     try:
-        FLORENCE_MODEL.config.use_cache = False
-
-        print(
-            "[Watermark] Florence config.use_cache=False"
-        )
+        if not hasattr(
+            FLORENCE_MODEL,
+            "_supports_sdpa",
+        ):
+            FLORENCE_MODEL._supports_sdpa = False
 
-    except Exception as e:
+    except Exception:
+        pass
 
-        print(
-            "[Watermark] Unable to set "
-            "config.use_cache=False:",
-            e,
-        )
+    # --------------------------------------------------------
+    # Print actual model dtype
+    # --------------------------------------------------------
 
     try:
-
-        FLORENCE_MODEL.generation_config.use_cache = False
+        actual_dtype = next(
+            FLORENCE_MODEL.parameters()
+        ).dtype
 
         print(
-            "[Watermark] Florence generation_config.use_cache=False"
+            f"[Florence] Actual model dtype: "
+            f"{actual_dtype}"
         )
 
     except Exception as e:
-
         print(
-            "[Watermark] Unable to set "
-            "generation_config.use_cache=False:",
-            e,
+            f"[Florence] Unable to determine model dtype: {e}"
         )
 
-    # --------------------------------------------------------
-    # Remove obsolete early_stopping warning
-    # --------------------------------------------------------
-
-    try:
-
-        if hasattr(
-            FLORENCE_MODEL.generation_config,
-            "early_stopping",
-        ):
-            FLORENCE_MODEL.generation_config.early_stopping = False
-
-    except Exception:
-        pass
-
     # ========================================================
-    # Install robust generation patch
-    # ========================================================
-
-    _patch_florence_generation(
-        FLORENCE_MODEL
-    )
-
-    # ========================================================
-    # LaMa
+    # Load LaMa
     # ========================================================
 
     lama_path = (
@@ -534,12 +329,11 @@ def _load_models():
 
         url = os.getenv(
             "LAMA_URL",
-            "https://github.com/Sanster/models/releases/"
-            "download/add_big_lama/big-lama.pt",
+            "https://github.com/Sanster/models/releases/download/add_big_lama/big-lama.pt",
         )
 
         print(
-            "[Watermark] Downloading LaMa model..."
+            "[LaMa] Downloading model..."
         )
 
         urllib.request.urlretrieve(
@@ -548,12 +342,12 @@ def _load_models():
         )
 
     print(
-        "[Watermark] Loading LaMa..."
+        f"[LaMa] Loading: {lama_path}"
     )
 
     LAMA_MODEL = torch.jit.load(
         str(lama_path),
-        map_location="cuda",
+        map_location=device,
     ).eval()
 
     print(
@@ -562,7 +356,7 @@ def _load_models():
 
 
 # ============================================================
-# Florence watermark detection
+# Florence detection
 # ============================================================
 
 def _boxes(
@@ -571,6 +365,13 @@ def _boxes(
 ):
     """
     Detect watermark bounding boxes using Florence-2.
+
+    IMPORTANT:
+    Florence model weights and pixel_values must use the
+    same floating-point dtype.
+
+    This function explicitly converts pixel_values to the
+    actual dtype of the Florence model.
     """
 
     _load_models()
@@ -580,6 +381,24 @@ def _boxes(
 
     image = image.convert("RGB")
 
+    device = _get_device()
+
+    # --------------------------------------------------------
+    # Get actual model dtype
+    # --------------------------------------------------------
+
+    try:
+        model_dtype = next(
+            FLORENCE_MODEL.parameters()
+        ).dtype
+
+    except Exception:
+        model_dtype = _get_model_dtype()
+
+    # --------------------------------------------------------
+    # Florence task
+    # --------------------------------------------------------
+
     task = (
         "<OPEN_VOCABULARY_DETECTION>"
     )
@@ -589,13 +408,11 @@ def _boxes(
         or "watermark"
     )
 
-    text = (
-        task + prompt
-    )
+    text = task + prompt
 
-    # ========================================================
+    # --------------------------------------------------------
     # Processor
-    # ========================================================
+    # --------------------------------------------------------
 
     inputs = FLORENCE_PROCESSOR(
         text=text,
@@ -603,78 +420,140 @@ def _boxes(
         return_tensors="pt",
     )
 
-    # ========================================================
-    # CUDA
-    # ========================================================
+    # --------------------------------------------------------
+    # Move tensors to CUDA
+    # --------------------------------------------------------
 
-    inputs = {
-        key: (
-            value.to("cuda")
-            if hasattr(value, "to")
-            else value
-        )
-        for key, value in inputs.items()
-    }
+    prepared = {}
 
-    # ========================================================
-    # Generation
-    # ========================================================
+    for key, value in inputs.items():
 
-    generation_kwargs = {
-        "max_new_tokens": 256,
-        "num_beams": 3,
+        if hasattr(value, "to"):
 
-        # Critical compatibility option.
-        "use_cache": False,
+            value = value.to(
+                device
+            )
 
-        # Avoid generation object overhead.
-        "return_dict_in_generate": False,
-    }
+        prepared[key] = value
 
     # --------------------------------------------------------
-    # Important:
+    # CRITICAL FIX
+    #
+    # Processor normally returns pixel_values as FP32.
+    #
+    # Florence-2-large can be loaded as BF16/FP16.
     #
-    # Florence's custom generate() expects:
+    # Conv2D requires:
     #
-    # input_ids
-    # pixel_values
+    # input dtype == weight/bias dtype
     #
-    # Do NOT pass None values.
+    # Therefore:
+    #
+    # pixel_values -> model dtype
     # --------------------------------------------------------
 
-    input_ids = inputs.get(
-        "input_ids"
-    )
-
-    pixel_values = inputs.get(
-        "pixel_values"
-    )
+    if "pixel_values" in prepared:
 
-    if input_ids is None:
-        raise RuntimeError(
-            "Florence processor did not return input_ids."
+        prepared[
+            "pixel_values"
+        ] = prepared[
+            "pixel_values"
+        ].to(
+            device=device,
+            dtype=model_dtype,
         )
 
-    if pixel_values is None:
-        raise RuntimeError(
-            "Florence processor did not return pixel_values."
+    # --------------------------------------------------------
+    # Token tensors must stay integer
+    # --------------------------------------------------------
+
+    if "input_ids" in prepared:
+
+        prepared[
+            "input_ids"
+        ] = prepared[
+            "input_ids"
+        ].long()
+
+    if "attention_mask" in prepared:
+
+        prepared[
+            "attention_mask"
+        ] = prepared[
+            "attention_mask"
+        ].long()
+
+    # --------------------------------------------------------
+    # Debug
+    # --------------------------------------------------------
+
+    pixel_dtype = None
+
+    if "pixel_values" in prepared:
+        pixel_dtype = (
+            prepared[
+                "pixel_values"
+            ].dtype
         )
 
-    # ========================================================
-    # Generate
-    # ========================================================
+    print(
+        "[Florence] "
+        f"model_dtype={model_dtype}, "
+        f"pixel_dtype={pixel_dtype}, "
+        f"device={device}"
+    )
+
+    # --------------------------------------------------------
+    # Generation
+    #
+    # Do NOT pass early_stopping.
+    #
+    # use_cache=False avoids the Florence remote-code
+    # past_key_values=None issue encountered previously.
+    # --------------------------------------------------------
+
+    generation_kwargs = {
+        "input_ids": prepared.get(
+            "input_ids"
+        ),
+
+        "pixel_values": prepared.get(
+            "pixel_values"
+        ),
+
+        "max_new_tokens": 1024,
+
+        "num_beams": 3,
+
+        "do_sample": False,
+
+        "use_cache": False,
+    }
+
+    # Remove None values
+
+    generation_kwargs = {
+        key: value
+        for key, value
+        in generation_kwargs.items()
+        if value is not None
+    }
+
+    # --------------------------------------------------------
+    # Florence inference
+    # --------------------------------------------------------
 
     with torch.inference_mode():
 
-        generated = FLORENCE_MODEL.generate(
-            input_ids=input_ids,
-            pixel_values=pixel_values,
-            **generation_kwargs,
+        generated = (
+            FLORENCE_MODEL.generate(
+                **generation_kwargs
+            )
         )
 
-    # ========================================================
+    # --------------------------------------------------------
     # Decode
-    # ========================================================
+    # --------------------------------------------------------
 
     decoded = (
         FLORENCE_PROCESSOR.batch_decode(
@@ -683,24 +562,24 @@ def _boxes(
         )[0]
     )
 
-    # ========================================================
-    # Post process
-    # ========================================================
+    print(
+        "[Florence] decoded:",
+        decoded[:500],
+    )
+
+    # --------------------------------------------------------
+    # Post processing
+    # --------------------------------------------------------
 
     parsed = (
-        FLORENCE_PROCESSOR.post_process_generation(
+        FLORENCE_PROCESSOR
+        .post_process_generation(
             decoded,
             task=task,
             image_size=image.size,
         )
     )
 
-    if not isinstance(
-        parsed,
-        dict,
-    ):
-        return []
-
     data = parsed.get(
         task,
         parsed,
@@ -720,9 +599,9 @@ def _boxes(
         or []
     )
 
-    # ========================================================
+    # --------------------------------------------------------
     # Validate boxes
-    # ========================================================
+    # --------------------------------------------------------
 
     valid_boxes = []
 
@@ -750,9 +629,7 @@ def _boxes(
         ):
             continue
 
-        # ----------------------------------------------------
         # Clamp
-        # ----------------------------------------------------
 
         x1 = max(
             0.0,
@@ -762,19 +639,19 @@ def _boxes(
             ),
         )
 
-        x2 = max(
+        y1 = max(
             0.0,
             min(
-                float(image.width),
-                x2,
+                float(image.height),
+                y1,
             ),
         )
 
-        y1 = max(
+        x2 = max(
             0.0,
             min(
-                float(image.height),
-                y1,
+                float(image.width),
+                x2,
             ),
         )
 
@@ -786,10 +663,10 @@ def _boxes(
             ),
         )
 
-        if x2 <= x1:
-            continue
-
-        if y2 <= y1:
+        if (
+            x2 <= x1
+            or y2 <= y1
+        ):
             continue
 
         valid_boxes.append(
@@ -802,15 +679,15 @@ def _boxes(
         )
 
     print(
-        "[Watermark] Florence detected "
-        f"{len(valid_boxes)} valid boxes."
+        f"[Florence] detected "
+        f"{len(valid_boxes)} bounding boxes"
     )
 
     return valid_boxes
 
 
 # ============================================================
-# Build mask
+# Create mask
 # ============================================================
 
 def _mask(
@@ -818,11 +695,6 @@ def _mask(
     boxes,
     max_pct,
 ):
-    """
-    Convert bounding boxes to an inpainting mask.
-    """
-
-    image = image.convert("RGB")
 
     mask = Image.new(
         "L",
@@ -839,80 +711,67 @@ def _mask(
         * image.height
     )
 
-    if total <= 0:
-        return mask
-
     for box in boxes:
 
         if len(box) != 4:
             continue
 
-        try:
-
-            x1, y1, x2, y2 = map(
-                float,
-                box,
-            )
-
-        except (
-            TypeError,
-            ValueError,
-        ):
-            continue
+        x1, y1, x2, y2 = map(
+            float,
+            box,
+        )
 
         x1 = max(
-            0.0,
+            0,
             min(
-                float(image.width),
+                image.width,
                 x1,
             ),
         )
 
         x2 = max(
-            0.0,
+            0,
             min(
-                float(image.width),
+                image.width,
                 x2,
             ),
         )
 
         y1 = max(
-            0.0,
+            0,
             min(
-                float(image.height),
+                image.height,
                 y1,
             ),
         )
 
         y2 = max(
-            0.0,
+            0,
             min(
-                float(image.height),
+                image.height,
                 y2,
             ),
         )
 
         area = (
             max(
-                0.0,
+                0,
                 x2 - x1,
             )
             *
             max(
-                0.0,
+                0,
                 y2 - y1,
             )
         )
 
-        area_pct = (
-            area
-            / total
-            * 100.0
-        )
-
-        if area_pct <= float(
-            max_pct
+        if (
+            total
+            and
+            area / total * 100
+            <= max_pct
         ):
+
             draw.rectangle(
                 (
                     x1,
@@ -923,12 +782,12 @@ def _mask(
                 fill=255,
             )
 
-    # --------------------------------------------------------
-    # Small expansion around watermark.
-    # --------------------------------------------------------
+    # Expand mask slightly
 
     return mask.filter(
-        ImageFilter.MaxFilter(7)
+        ImageFilter.MaxFilter(
+            7
+        )
     )
 
 
@@ -942,15 +801,16 @@ def _lama(
     boxes=None,
 ):
     """
-    Run LaMa inpainting.
+    Run LaMa on watermark regions.
 
-    When boxes exist, process local regions instead of
-    the entire frame. This is significantly cheaper for
-    video processing.
+    For video, process only the detected regions instead of
+    the complete 1080p/4K frame.
     """
 
     _load_models()
 
+    device = _get_device()
+
     image = image.convert(
         "RGB"
     )
@@ -960,7 +820,7 @@ def _lama(
     )
 
     # ========================================================
-    # Region-based processing
+    # Region-based inpainting
     # ========================================================
 
     if boxes:
@@ -970,7 +830,9 @@ def _lama(
         pad = max(
             16,
             int(
-                min(image.size)
+                min(
+                    image.size
+                )
                 * 0.02
             ),
         )
@@ -980,18 +842,10 @@ def _lama(
             if len(box) != 4:
                 continue
 
-            try:
-
-                x1, y1, x2, y2 = map(
-                    float,
-                    box,
-                )
-
-            except (
-                TypeError,
-                ValueError,
-            ):
-                continue
+            x1, y1, x2, y2 = map(
+                float,
+                box,
+            )
 
             left = max(
                 0,
@@ -1013,10 +867,10 @@ def _lama(
                 int(y2) + pad,
             )
 
-            if right <= left:
-                continue
-
-            if bottom <= top:
+            if (
+                right <= left
+                or bottom <= top
+            ):
                 continue
 
             crop = image.crop(
@@ -1037,10 +891,6 @@ def _lama(
                 )
             )
 
-            # ------------------------------------------------
-            # Image tensor
-            # ------------------------------------------------
-
             crop_np = (
                 np.asarray(
                     crop,
@@ -1067,8 +917,11 @@ def _lama(
                     1,
                 )
                 .unsqueeze(0)
-                .cuda()
-                * 2.0
+                .to(device)
+            )
+
+            image_t = (
+                image_t * 2.0
                 - 1.0
             )
 
@@ -1078,17 +931,13 @@ def _lama(
                 )
                 .unsqueeze(0)
                 .unsqueeze(0)
-                .cuda()
+                .to(device)
             )
 
             mask_t = (
                 mask_t > 0.5
             ).float()
 
-            # ------------------------------------------------
-            # LaMa
-            # ------------------------------------------------
-
             with torch.inference_mode():
 
                 out = LAMA_MODEL(
@@ -1127,7 +976,8 @@ def _lama(
 
             restored = Image.fromarray(
                 (
-                    out * 255
+                    out
+                    * 255
                 )
                 .clip(
                     0,
@@ -1149,7 +999,7 @@ def _lama(
         return result
 
     # ========================================================
-    # Full-frame fallback
+    # Full image fallback
     # ========================================================
 
     image_np = (
@@ -1178,8 +1028,11 @@ def _lama(
             1,
         )
         .unsqueeze(0)
-        .cuda()
-        * 2.0
+        .to(device)
+    )
+
+    image_t = (
+        image_t * 2.0
         - 1.0
     )
 
@@ -1189,7 +1042,7 @@ def _lama(
         )
         .unsqueeze(0)
         .unsqueeze(0)
-        .cuda()
+        .to(device)
     )
 
     mask_t = (
@@ -1234,7 +1087,8 @@ def _lama(
 
     return Image.fromarray(
         (
-            result * 255
+            result
+            * 255
         )
         .clip(
             0,
@@ -1255,25 +1109,22 @@ def preview(
     max_bbox_percent=30,
     detection_prompt="watermark",
 ):
-    """
-    Detect watermark and draw bounding boxes.
-    """
 
     if image is None:
         raise ValueError(
             "Please upload an image."
         )
 
-    image = image.convert(
-        "RGB"
-    )
-
     boxes = _boxes(
         image,
         detection_prompt,
     )
 
-    out = image.copy()
+    out = (
+        image
+        .convert("RGB")
+        .copy()
+    )
 
     draw = ImageDraw.Draw(
         out
@@ -1281,19 +1132,18 @@ def preview(
 
     for box in boxes:
 
-        if len(box) != 4:
-            continue
+        if len(box) == 4:
 
-        draw.rectangle(
-            tuple(
-                map(
-                    float,
-                    box,
-                )
-            ),
-            outline="red",
-            width=4,
-        )
+            draw.rectangle(
+                tuple(
+                    map(
+                        float,
+                        box,
+                    )
+                ),
+                outline="red",
+                width=4,
+            )
 
     return out
 
@@ -1308,58 +1158,37 @@ def remove_image(
     detection_prompt="watermark",
     transparent=False,
 ):
-    """
-    Detect and remove watermark from image.
-    """
 
     if image is None:
         raise ValueError(
             "Please upload an image."
         )
 
-    image = image.convert(
-        "RGB"
-    )
-
-    # --------------------------------------------------------
-    # Detect
-    # --------------------------------------------------------
-
     boxes = _boxes(
         image,
         detection_prompt,
     )
 
-    # --------------------------------------------------------
-    # Mask
-    # --------------------------------------------------------
-
     mask = _mask(
         image,
         boxes,
         max_bbox_percent,
     )
 
-    # --------------------------------------------------------
-    # Nothing detected
-    # --------------------------------------------------------
-
     if not boxes:
 
-        result = image.convert(
+        output_mode = (
             "RGBA"
             if transparent
             else "RGB"
         )
 
-        return result, mask
-
-    # --------------------------------------------------------
-    # LaMa
-    #
-    # IMPORTANT:
-    # Pass boxes here so LaMa uses region processing.
-    # --------------------------------------------------------
+        return (
+            image.convert(
+                output_mode
+            ),
+            mask,
+        )
 
     result = _lama(
         image,
@@ -1367,10 +1196,6 @@ def remove_image(
         boxes=boxes,
     )
 
-    # --------------------------------------------------------
-    # Transparent output
-    # --------------------------------------------------------
-
     if transparent:
 
         result = result.convert(
@@ -1379,8 +1204,12 @@ def remove_image(
 
         result.putalpha(
             mask.point(
-                lambda p: 255 - p
+                lambda p:
+                    255 - p
             )
         )
 
-    return result, mask
+    return (
+        result,
+        mask,
+    )