Download watermark/engine.py from bened/Z-Image: direct link, hf CLI and curl.
- Browser
- Download file 12.5 kB
-
https://huggingface.co/spaces/bened/Z-Image/resolve/main/watermark/engine.py
- Command line
-
hf download hf://spaces/bened/Z-Image/watermark/engine.py
-
curl -L -o engine.py https://huggingface.co/spaces/bened/Z-Image/resolve/main/watermark/engine.py
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 | |