Z-Image / watermark /engine.py
bened's picture
Upload 2 files
eb7f2ce verified
Raw History Blame Contribute Delete
12.5 kB
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