feat: watermark assets + skip R2 upload for web-UI generations

Add a brand watermark (logo + domain) compositor applied to generated images
when the caller passes a valid JSON spec via the wm cookie, and stop uploading
ad-hoc Gradio web-UI generations to R2 (only API/generator calls, detected via
the uid cookie / /call REST path, populate the asset store).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

app.py

@@ -48,6 +48,7 @@ def _progress(stage, p, step=0, total=0, label=""):
     }
 
 import r2_uploader
+import watermark as watermark_module
 
 # Per-Space namespace embedded in every uploaded object key. Deliberately opaque
 # (not the readable Space name) but stable so the owner can tell assets apart.
@@ -1194,6 +1195,20 @@ def generate_and_upload(
         else:
             _, image, used, moderation, effective_prompt = ev
 
+    # Stamp the brand watermark (logo + domain) onto the asset when the caller
+    # supplied a valid ``wm`` cookie spec — applied after moderation so it never
+    # affects the safety rating, and to the same image we display and upload.
+    watermark = r2_uploader.watermark_from_request(request)
+    if watermark_module.is_valid(watermark):
+        image = watermark_module.apply_to_image(image, watermark)
+
+    # Direct Gradio web-UI generations are not uploaded to R2 — only API calls
+    # (the generator, which forwards a uid cookie) populate the asset store.
+    if not r2_uploader.request_is_api(request):
+        status = {"r2_skipped": "web-ui generation (not uploaded)", "moderation": moderation}
+        yield image, used, status, _progress("done", 1.0, label="Done")
+        return
+
     yield None, used, None, _progress("image", 0.97, label="Uploading")
     uid = r2_uploader.uid_from_request(request)
     buf = io.BytesIO()

r2_uploader.py

@@ -72,6 +72,74 @@ def uid_from_request(request, cookie_name: str = "uid") -> str:
     return ""
 
 
+def watermark_from_request(request, cookie_name: str = "wm") -> str:
+    """Best-effort read of the watermark JSON spec from the ``wm`` cookie.
+
+    The generator forwards the brand watermark spec (a JSON string such as
+    ``{"logo":"https://.../logo.png","domain":"nsfwailab.com"}``) as the ``wm``
+    cookie on the Space call — the same channel as the ``uid`` cookie — so no
+    Gradio input arity changes (which would break older callers). The value is
+    percent-encoded (encodeURIComponent) by the caller and decoded here. Returns
+    "" when absent; the web UI never sets it, so UI generations stay un-marked.
+    """
+    if request is None:
+        return ""
+    from urllib.parse import unquote
+
+    try:
+        raw = ""
+        cookies = getattr(request, "cookies", None)
+        if isinstance(cookies, dict) and cookies.get(cookie_name):
+            raw = str(cookies[cookie_name])
+        else:
+            headers = getattr(request, "headers", None)
+            header_cookie = ""
+            if headers is not None:
+                header_cookie = (headers.get("cookie") if hasattr(headers, "get") else "") or ""
+            for part in header_cookie.split(";"):
+                k, _, v = part.strip().partition("=")
+                if k == cookie_name and v:
+                    raw = v
+                    break
+        if not raw:
+            return ""
+        decoded = unquote(raw)
+        # Tolerate accidental double-encoding (value still looks percent-encoded).
+        if "%7B" in decoded or "%22" in decoded:
+            decoded = unquote(decoded)
+        return decoded
+    except Exception:  # noqa: BLE001 - watermark is best-effort, never fatal
+        return ""
+
+def request_is_api(request) -> bool:
+    """True when a generation call arrived through the REST API (the generator),
+    False for a direct Gradio web-UI interaction.
+
+    Used to skip the R2 upload for ad-hoc web-UI generations: the generator
+    always forwards the requester's ``uid`` cookie and reaches the Space through
+    the ``/gradio_api/call/<api>`` REST endpoint, whereas a browser on the public
+    Gradio UI does neither (it uses the queue path with no uid cookie).
+    """
+    if request is None:
+        return False
+    if uid_from_request(request):
+        return True
+    try:
+        path = ""
+        url = getattr(request, "url", None)
+        if url is not None:
+            path = getattr(url, "path", "") or str(url)
+        if not path:
+            raw = getattr(request, "request", None)
+            url = getattr(raw, "url", None) if raw is not None else None
+            path = getattr(url, "path", "") if url is not None else ""
+        if "/call/" in path or path.endswith("/call"):
+            return True
+    except Exception:  # noqa: BLE001 - detection is best-effort
+        pass
+    return False
+
+
 def _client(cfg: dict):
     import boto3
     from botocore.config import Config

watermark.py

@@ -0,0 +1,311 @@
+"""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