Delete watermark/engine_v3.py

watermark/engine_v3.py

@@ -1,455 +0,0 @@
-from pathlib import Path
-import os
-import urllib.request
-
-import numpy as np
-import torch
-from PIL import Image, ImageDraw, ImageFilter
-
-MODEL_CACHE = None
-FLORENCE_MODEL = None
-FLORENCE_PROCESSOR = None
-LAMA_MODEL = None
-
-
-def _get_device():
-    return torch.device("cuda" if torch.cuda.is_available() else "cpu")
-
-
-def _get_model_dtype():
-    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
-
-
-def _get_writable_cache():
-    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:
-        try:
-            p.mkdir(parents=True, exist_ok=True)
-            test = p / ".write_test"
-            test.write_text("ok", encoding="utf-8")
-            try:
-                test.unlink()
-            except Exception:
-                pass
-            print(f"[Watermark] Cache directory: {p}")
-            return p
-        except (PermissionError, OSError):
-            pass
-    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
-
-
-def _patch_florence_compatibility():
-    """Compatibility patches for Florence-2 remote code + newer Transformers."""
-    try:
-        from transformers import PretrainedConfig, PreTrainedModel
-
-        if not hasattr(PretrainedConfig, "forced_bos_token_id"):
-            PretrainedConfig.forced_bos_token_id = None
-            print("[Florence] Applied forced_bos_token_id compatibility patch.")
-
-        if not hasattr(PreTrainedModel, "_supports_sdpa"):
-            PreTrainedModel._supports_sdpa = False
-            print("[Florence] Applied _supports_sdpa compatibility patch.")
-    except Exception as e:
-        print(f"[Florence] Compatibility patch skipped: {e}")
-
-
-def _load_models():
-    global FLORENCE_MODEL, FLORENCE_PROCESSOR, LAMA_MODEL
-
-    if (
-        FLORENCE_MODEL is not None
-        and FLORENCE_PROCESSOR is not None
-        and LAMA_MODEL is not None
-    ):
-        return
-
-    _patch_florence_compatibility()
-
-    from transformers import AutoProcessor, AutoModelForCausalLM
-
-    cache = _cache_dir()
-    model_id = os.getenv(
-        "FLORENCE_MODEL_ID",
-        "microsoft/Florence-2-large",
-    )
-    device = _get_device()
-    dtype = _get_model_dtype()
-
-    print(f"[Florence] Loading model: {model_id}")
-    print(f"[Florence] device={device}, dtype={dtype}")
-
-    FLORENCE_PROCESSOR = AutoProcessor.from_pretrained(
-        model_id,
-        cache_dir=str(cache / "hf"),
-        trust_remote_code=True,
-    )
-
-    model_kwargs = {
-        "cache_dir": str(cache / "hf"),
-        "trust_remote_code": True,
-        "torch_dtype": dtype,
-    }
-
-    # Force eager attention so Transformers does not try to negotiate
-    # SDPA with older Florence remote-code revisions.
-    try:
-        model_kwargs["attn_implementation"] = "eager"
-    except Exception:
-        pass
-
-    FLORENCE_MODEL = AutoModelForCausalLM.from_pretrained(
-        model_id,
-        **model_kwargs,
-    ).to(device).eval()
-
-    # Patch the actual custom model instance as well.
-    try:
-        if not hasattr(FLORENCE_MODEL, "_supports_sdpa"):
-            FLORENCE_MODEL._supports_sdpa = False
-    except Exception:
-        pass
-
-    try:
-        actual_dtype = next(FLORENCE_MODEL.parameters()).dtype
-        print(f"[Florence] model dtype={actual_dtype}")
-    except Exception:
-        pass
-
-    lama_path = cache / "big-lama.pt"
-    if not lama_path.exists():
-        url = os.getenv(
-            "LAMA_URL",
-            "https://github.com/Sanster/models/releases/download/add_big_lama/big-lama.pt",
-        )
-        print("[LaMa] Downloading model...")
-        urllib.request.urlretrieve(url, lama_path)
-
-    print("[LaMa] Loading model...")
-    LAMA_MODEL = torch.jit.load(
-        str(lama_path),
-        map_location=device,
-    ).eval()
-
-    print("[Watermark] All watermark models loaded.")
-
-
-def _boxes(image, prompt):
-    """Detect watermark bounding boxes using Florence-2."""
-    _load_models()
-
-    if image is None:
-        return []
-
-    image = image.convert("RGB")
-    device = _get_device()
-
-    try:
-        model_dtype = next(FLORENCE_MODEL.parameters()).dtype
-    except Exception:
-        model_dtype = _get_model_dtype()
-
-    task = "<OPEN_VOCABULARY_DETECTION>"
-    text = task + (prompt or "watermark")
-
-    inputs = FLORENCE_PROCESSOR(
-        text=text,
-        images=image,
-        return_tensors="pt",
-    )
-
-    prepared = {}
-    for key, value in inputs.items():
-        if hasattr(value, "to"):
-            value = value.to(device)
-        prepared[key] = value
-
-    # Florence vision tower Conv2D requires input dtype to match model weights.
-    if "pixel_values" in prepared:
-        prepared["pixel_values"] = prepared["pixel_values"].to(
-            device=device,
-            dtype=model_dtype,
-        )
-
-    if "input_ids" in prepared:
-        prepared["input_ids"] = prepared["input_ids"].long()
-    if "attention_mask" in prepared:
-        prepared["attention_mask"] = prepared["attention_mask"].long()
-
-    print(
-        f"[Florence] model_dtype={model_dtype}, "
-        f"pixel_dtype={prepared.get('pixel_values').dtype if 'pixel_values' in prepared else None}, "
-        f"device={device}"
-    )
-
-    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,
-    }
-    generation_kwargs = {
-        k: v for k, v in generation_kwargs.items() if v is not None
-    }
-
-    with torch.inference_mode():
-        generated = FLORENCE_MODEL.generate(**generation_kwargs)
-
-    decoded = FLORENCE_PROCESSOR.batch_decode(
-        generated,
-        skip_special_tokens=False,
-    )[0]
-
-    print("[Florence] decoded:", decoded[:500])
-
-    parsed = FLORENCE_PROCESSOR.post_process_generation(
-        decoded,
-        task=task,
-        image_size=image.size,
-    )
-    data = parsed.get(task, parsed)
-
-    if not isinstance(data, dict):
-        return []
-
-    raw_boxes = data.get("bboxes", []) or []
-    valid_boxes = []
-
-    for box in raw_boxes:
-        if not isinstance(box, (list, tuple)) or len(box) != 4:
-            continue
-        try:
-            x1, y1, x2, y2 = map(float, box)
-        except (TypeError, ValueError):
-            continue
-
-        x1 = max(0.0, min(float(image.width), x1))
-        y1 = max(0.0, min(float(image.height), y1))
-        x2 = max(0.0, min(float(image.width), x2))
-        y2 = max(0.0, min(float(image.height), y2))
-
-        if x2 <= x1 or y2 <= y1:
-            continue
-
-        valid_boxes.append([x1, y1, x2, y2])
-
-    print(f"[Florence] detected {len(valid_boxes)} bounding boxes")
-    return valid_boxes
-
-
-def _mask(image, boxes, max_pct):
-    mask = Image.new("L", image.size, 0)
-    draw = ImageDraw.Draw(mask)
-    total = image.width * image.height
-
-    for box in boxes or []:
-        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)
-
-    # Small expansion around the detected watermark.
-    return mask.filter(ImageFilter.MaxFilter(7))
-
-
-def _to_lama_tensors(image, mask, device, multiple=8):
-    """Convert an image/mask crop to LaMa tensors with safe dimensions."""
-    image = image.convert("RGB")
-    mask = mask.convert("L")
-
-    w, h = image.size
-    pad_w = (multiple - (w % multiple)) % multiple
-    pad_h = (multiple - (h % multiple)) % multiple
-
-    image_np = np.asarray(image, dtype=np.float32) / 255.0
-    mask_np = np.asarray(mask, dtype=np.float32) / 255.0
-
-    if pad_w or pad_h:
-        # Edge padding keeps the context natural and avoids black borders.
-        image_np = np.pad(
-            image_np,
-            ((0, pad_h), (0, pad_w), (0, 0)),
-            mode="edge",
-        )
-        mask_np = np.pad(
-            mask_np,
-            ((0, pad_h), (0, pad_w)),
-            mode="constant",
-            constant_values=0,
-        )
-
-    image_t = torch.from_numpy(image_np).permute(2, 0, 1).unsqueeze(0).to(device)
-    image_t = image_t * 2.0 - 1.0
-
-    mask_t = torch.from_numpy(mask_np).unsqueeze(0).unsqueeze(0).to(device)
-    mask_t = (mask_t > 0.5).float()
-
-    return image_t, mask_t, (w, h)
-
-
-def _run_lama(image, mask):
-    """Run LaMa on a crop and return a crop-sized PIL image."""
-    device = _get_device()
-    image_t, mask_t, original_size = _to_lama_tensors(
-        image,
-        mask,
-        device,
-        multiple=8,
-    )
-
-    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)
-    )
-
-    # Remove any safety padding added for LaMa.
-    restored = restored.crop(
-        (0, 0, original_size[0], original_size[1])
-    )
-    return restored
-
-
-def _lama(image, mask, boxes=None):
-    """
-    Run LaMa only on watermark regions.
-
-    Important V2/V3 fixes:
-      1. Pad crop dimensions to a multiple of 8 to avoid LaMa tensor-size
-         mismatches such as 245 vs 248.
-      2. Composite the LaMa result through the actual mask. The old code pasted
-         the entire restored crop back into the frame, which could create a
-         bright/white rectangular patch around the watermark.
-      3. Keep the original pixels outside the mask untouched.
-    """
-    _load_models()
-
-    image = image.convert("RGB")
-    mask = mask.convert("L")
-
-    if not boxes:
-        restored = _run_lama(image, mask)
-        blend_mask = mask.filter(ImageFilter.GaussianBlur(1.0))
-        return Image.composite(restored, image, blend_mask)
-
-    result = image.copy()
-    pad = max(16, int(min(image.size) * 0.02))
-
-    for box in boxes:
-        if not isinstance(box, (list, tuple)) or 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:
-            continue
-
-        crop = result.crop((left, top, right, bottom))
-        crop_mask = mask.crop((left, top, right, bottom))
-
-        if not np.asarray(crop_mask).any():
-            continue
-
-        restored = _run_lama(crop, crop_mask)
-
-        # Feather only the transition edge. Do NOT paste the entire LaMa crop.
-        blend_mask = crop_mask.filter(ImageFilter.GaussianBlur(1.2))
-        blended = Image.composite(
-            restored,
-            crop,
-            blend_mask,
-        )
-
-        result.paste(
-            blended,
-            (left, top),
-        )
-
-    return result
-
-
-def preview(image, max_bbox_percent=30, detection_prompt="watermark"):
-    if image is None:
-        raise ValueError("Please upload an image.")
-
-    boxes = _boxes(image, detection_prompt)
-    out = image.convert("RGB").copy()
-    draw = ImageDraw.Draw(out)
-
-    for box in boxes:
-        if len(box) == 4:
-            draw.rectangle(
-                tuple(map(float, box)),
-                outline="red",
-                width=4,
-            )
-
-    return out
-
-
-def remove_image(
-    image,
-    max_bbox_percent=30,
-    detection_prompt="watermark",
-    transparent=False,
-):
-    if image is None:
-        raise ValueError("Please upload an image.")
-
-    boxes = _boxes(image, detection_prompt)
-    mask = _mask(image, boxes, max_bbox_percent)
-
-    if not boxes:
-        return image.convert("RGBA" if transparent else "RGB"), mask
-
-    result = _lama(image, mask, boxes=boxes)
-
-    if transparent:
-        result = result.convert("RGBA")
-        result.putalpha(mask.point(lambda p: 255 - p))
-
-    return result, mask