Spaces:
Running on Zero
Running on Zero
| """Brand watermark overlay for generated images and videos. | |
| A watermark is requested by passing a JSON spec (a string or an already-parsed | |
| dict) to a generation endpoint via the ``watermark`` parameter. When the spec is | |
| valid, a small badge — the brand logo with the site domain rendered beneath it — | |
| is composited onto the generated image (or every frame of the generated video, | |
| via a single ffmpeg overlay pass). An empty or invalid spec is a no-op, so any | |
| caller that omits the parameter is completely unaffected. | |
| Spec fields (JSON object; at least one of ``logo`` / ``domain`` must be present | |
| for the spec to count as "valid"): | |
| logo str logo image as an http(s) URL, a ``data:`` URI, or raw | |
| base64. Optional — a domain-only watermark is allowed. | |
| domain str text drawn under the logo, e.g. "nsfwailab.com". Optional. | |
| position str bottom-right (default), bottom-left, top-right, top-left, | |
| bottom-center or top-center. | |
| scale float badge width as a fraction of the media width (default 0.18, | |
| clamped to 0.05..0.6). | |
| opacity float overall badge opacity 0..1 (default 0.9). | |
| margin float badge margin from the edge as a fraction of media width | |
| (default 0.03). | |
| Every public entry point is wrapped so a watermarking failure degrades to | |
| "asset returned un-watermarked" rather than breaking generation. | |
| """ | |
| from __future__ import annotations | |
| import base64 | |
| import json | |
| import os | |
| import subprocess | |
| import tempfile | |
| import urllib.request | |
| _VIDEO_REF_W = 1024 # fallback media width when a video's size can't be probed. | |
| # Cache decoded logos so a per-frame / per-call render doesn't refetch or | |
| # re-decode the same image. Keyed on the raw logo string. | |
| _logo_cache: dict[str, object] = {} | |
| _POSITIONS = { | |
| "bottom-right", "bottom-left", "top-right", "top-left", | |
| "bottom-center", "top-center", | |
| } | |
| def parse_spec(spec) -> dict | None: | |
| """Validate/normalise a watermark spec; return a clean dict or None. | |
| Accepts a JSON string or a dict. Returns None (a no-op signal) when the spec | |
| is empty, unparseable, or carries neither a logo nor a domain. | |
| """ | |
| if not spec: | |
| return None | |
| data = spec | |
| if isinstance(spec, str): | |
| s = spec.strip() | |
| if not s: | |
| return None | |
| try: | |
| data = json.loads(s) | |
| except (ValueError, TypeError): | |
| return None | |
| if not isinstance(data, dict): | |
| return None | |
| logo = data.get("logo") or data.get("logo_url") or "" | |
| domain = (data.get("domain") or data.get("text") or "").strip() | |
| if not (logo or domain): | |
| return None | |
| def _f(key, default, lo, hi): | |
| try: | |
| return max(lo, min(hi, float(data.get(key, default)))) | |
| except (TypeError, ValueError): | |
| return default | |
| position = str(data.get("position", "bottom-right")).strip().lower() | |
| if position not in _POSITIONS: | |
| position = "bottom-right" | |
| return { | |
| "logo": logo if isinstance(logo, str) else "", | |
| "domain": domain, | |
| "position": position, | |
| "scale": _f("scale", 0.18, 0.05, 0.6), | |
| "opacity": _f("opacity", 0.9, 0.1, 1.0), | |
| "margin": _f("margin", 0.03, 0.0, 0.2), | |
| } | |
| def is_valid(spec) -> bool: | |
| """True when ``spec`` would produce a watermark.""" | |
| return parse_spec(spec) is not None | |
| def _load_logo(logo: str): | |
| """Decode a logo string (URL / data-URI / base64) into an RGBA PIL image.""" | |
| if not logo: | |
| return None | |
| if logo in _logo_cache: | |
| return _logo_cache[logo] | |
| from PIL import Image | |
| import io | |
| img = None | |
| try: | |
| raw = None | |
| if logo.startswith("data:"): | |
| _, _, b64 = logo.partition(",") | |
| raw = base64.b64decode(b64) | |
| elif logo.startswith("http://") or logo.startswith("https://"): | |
| req = urllib.request.Request(logo, headers={"User-Agent": "watermark/1.0"}) | |
| with urllib.request.urlopen(req, timeout=10) as resp: | |
| raw = resp.read() | |
| else: | |
| raw = base64.b64decode(logo) | |
| if raw: | |
| img = Image.open(io.BytesIO(raw)).convert("RGBA") | |
| except Exception as exc: # noqa: BLE001 - logo is best-effort | |
| print(f"[watermark] logo decode failed: {exc}") | |
| img = None | |
| _logo_cache[logo] = img | |
| return img | |
| _FONT_CANDIDATES = [ | |
| "/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf", | |
| "/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf", | |
| "/usr/share/fonts/truetype/liberation/LiberationSans-Bold.ttf", | |
| "DejaVuSans-Bold.ttf", | |
| ] | |
| def _font(size: int): | |
| from PIL import ImageFont | |
| for path in _FONT_CANDIDATES: | |
| try: | |
| return ImageFont.truetype(path, size) | |
| except Exception: # noqa: BLE001 | |
| continue | |
| return ImageFont.load_default() | |
| def _render_badge(cfg: dict, media_w: int): | |
| """Build the watermark badge (logo + domain on a translucent plate).""" | |
| from PIL import Image, ImageDraw | |
| badge_w = max(72, int(media_w * cfg["scale"])) | |
| pad = max(6, badge_w // 16) | |
| logo = _load_logo(cfg["logo"]) | |
| logo_img = None | |
| logo_h = 0 | |
| if logo is not None and logo.width > 0: | |
| lw = max(1, badge_w - 2 * pad) | |
| lh = max(1, int(logo.height * lw / logo.width)) | |
| logo_img = logo.resize((lw, lh), Image.LANCZOS) | |
| logo_h = lh | |
| domain = cfg["domain"] | |
| font = None | |
| text_w = text_h = 0 | |
| bbox = (0, 0, 0, 0) | |
| if domain: | |
| measure = ImageDraw.Draw(Image.new("RGBA", (1, 1))) | |
| fs = max(11, int(badge_w * 0.14)) | |
| font = _font(fs) | |
| bbox = measure.textbbox((0, 0), domain, font=font) | |
| text_w, text_h = bbox[2] - bbox[0], bbox[3] - bbox[1] | |
| max_tw = badge_w - 2 * pad | |
| if text_w > max_tw and text_w > 0: | |
| fs = max(8, int(fs * max_tw / text_w)) | |
| font = _font(fs) | |
| bbox = measure.textbbox((0, 0), domain, font=font) | |
| text_w, text_h = bbox[2] - bbox[0], bbox[3] - bbox[1] | |
| gap = pad if (logo_img is not None and domain) else 0 | |
| badge_h = logo_h + gap + text_h + 2 * pad | |
| badge = Image.new("RGBA", (badge_w, badge_h), (0, 0, 0, 0)) | |
| draw = ImageDraw.Draw(badge) | |
| radius = max(8, badge_w // 12) | |
| draw.rounded_rectangle([0, 0, badge_w - 1, badge_h - 1], radius=radius, | |
| fill=(0, 0, 0, 110)) | |
| y = pad | |
| if logo_img is not None: | |
| badge.alpha_composite(logo_img, ((badge_w - logo_img.width) // 2, y)) | |
| y += logo_h + gap | |
| if domain: | |
| tx = (badge_w - text_w) // 2 - bbox[0] | |
| draw.text((tx, y - bbox[1]), domain, font=font, fill=(255, 255, 255, 235)) | |
| opacity = cfg["opacity"] | |
| if opacity < 1.0: | |
| alpha = badge.split()[3].point(lambda v: int(v * opacity)) | |
| badge.putalpha(alpha) | |
| return badge | |
| def _offset(position: str, media_w: int, media_h: int, bw: int, bh: int, margin: int): | |
| """Top-left pixel offset for the badge given a named position.""" | |
| if "right" in position: | |
| x = media_w - bw - margin | |
| elif "left" in position: | |
| x = margin | |
| else: # center | |
| x = (media_w - bw) // 2 | |
| y = margin if position.startswith("top") else media_h - bh - margin | |
| return max(0, x), max(0, y) | |
| def apply_to_image(img, spec): | |
| """Return ``img`` with the watermark composited on, or unchanged on no-op.""" | |
| cfg = parse_spec(spec) | |
| if cfg is None: | |
| return img | |
| try: | |
| base = img.convert("RGBA") | |
| badge = _render_badge(cfg, base.width) | |
| margin = int(cfg["margin"] * base.width) | |
| x, y = _offset(cfg["position"], base.width, base.height, | |
| badge.width, badge.height, margin) | |
| base.alpha_composite(badge, (x, y)) | |
| return base.convert("RGB") | |
| except Exception as exc: # noqa: BLE001 - never break generation | |
| print(f"[watermark] image overlay failed: {exc}") | |
| return img | |
| def _video_dims(path: str): | |
| try: | |
| import cv2 | |
| cap = cv2.VideoCapture(path) | |
| w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) | |
| h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) | |
| cap.release() | |
| if w > 0 and h > 0: | |
| return w, h | |
| except Exception: # noqa: BLE001 | |
| pass | |
| # Fallback when OpenCV is unavailable: parse the resolution out of ffmpeg's | |
| # probe output so the badge is still sized/placed against the real frame. | |
| try: | |
| import re | |
| import imageio_ffmpeg | |
| ff = imageio_ffmpeg.get_ffmpeg_exe() | |
| err = subprocess.run([ff, "-i", path], capture_output=True, text=True).stderr | |
| m = re.search(r"Video:.*?(\d{2,5})x(\d{2,5})", err) | |
| if m: | |
| return int(m.group(1)), int(m.group(2)) | |
| except Exception: # noqa: BLE001 | |
| pass | |
| return None | |
| def apply_to_video(in_path: str, spec, out_path: str | None = None) -> str: | |
| """Overlay the watermark on a video via one ffmpeg pass. | |
| Returns the path to the watermarked file, or the original ``in_path`` when | |
| the spec is a no-op or anything fails (so upload still proceeds). | |
| """ | |
| cfg = parse_spec(spec) | |
| if cfg is None: | |
| return in_path | |
| badge_png = None | |
| try: | |
| import imageio_ffmpeg | |
| ffmpeg = imageio_ffmpeg.get_ffmpeg_exe() | |
| dims = _video_dims(in_path) | |
| media_w, media_h = dims if dims else (_VIDEO_REF_W, _VIDEO_REF_W) | |
| badge = _render_badge(cfg, media_w) | |
| margin = int(cfg["margin"] * media_w) | |
| x, y = _offset(cfg["position"], media_w, media_h, | |
| badge.width, badge.height, margin) | |
| fd, badge_png = tempfile.mkstemp(suffix=".png") | |
| os.close(fd) | |
| badge.save(badge_png) | |
| if out_path is None: | |
| fd, out_path = tempfile.mkstemp(suffix=".mp4") | |
| os.close(fd) | |
| cmd = [ | |
| ffmpeg, "-y", "-i", in_path, "-i", badge_png, | |
| "-filter_complex", f"overlay={x}:{y}:format=auto", | |
| "-c:a", "copy", "-c:v", "libx264", "-pix_fmt", "yuv420p", | |
| "-crf", "18", "-preset", "veryfast", "-movflags", "+faststart", | |
| out_path, | |
| ] | |
| subprocess.run(cmd, check=True, capture_output=True) | |
| return out_path | |
| except Exception as exc: # noqa: BLE001 - never break generation | |
| detail = exc.stderr.decode("utf-8", "ignore")[-400:] if isinstance(exc, subprocess.CalledProcessError) else exc | |
| print(f"[watermark] video overlay failed: {detail}") | |
| return in_path | |
| finally: | |
| if badge_png and os.path.exists(badge_png): | |
| try: | |
| os.unlink(badge_png) | |
| except OSError: | |
| pass | |